ENH: Update with TypedDict

Update using TypedDict for state

6 files changed, 53 insertions, 28 deletions

diff --git a/numpy/random/_generator.pyi b/numpy/random/_generator.pyi
index 22860b5ea..648e1d29d 100644
--- a/numpy/random/_generator.pyi
+++ b/numpy/random/_generator.pyi
@@ -3,7 +3,15 @@ from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union, overlo
 
 from numpy import dtype, float32, float64, int64, integer, ndarray
 from numpy.random import BitGenerator
-from numpy.typing import ArrayLike, DTypeLike, _ArrayLikeFloat_co, _ArrayLikeInt_co, _ShapeLike, _DoubleCodes, _SingleCodes
+from numpy.typing import (
+    ArrayLike,
+    DTypeLike,
+    _ArrayLikeFloat_co,
+    _ArrayLikeInt_co,
+    _DoubleCodes,
+    _ShapeLike,
+    _SingleCodes,
+)
 
 if sys.version_info >= (3, 8):
     from typing import Literal
@@ -72,15 +80,14 @@ class Generator:
         size: Union[None, _ShapeLike] = ...,
     ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
     @overload
-    def standard_normal(
+    def standard_normal(  # type: ignore[misc]
         self,
         size: None = ...,
         dtype: DTypeLike = ...,
         out: None = ...,
     ) -> float: ...
-    # TODO: How to literal dtype?
     @overload
-    def standard_normal(
+    def standard_normal(  # type: ignore[misc]
         self,
         size: _ShapeLike = ...,
         dtype: DTypeLike = ...,
diff --git a/numpy/random/_mt19937.pyi b/numpy/random/_mt19937.pyi
index 91eccda38..edd433598 100644
--- a/numpy/random/_mt19937.pyi
+++ b/numpy/random/_mt19937.pyi
@@ -1,10 +1,16 @@
-from typing import Any, Dict, Union
+from typing import Any, Dict, TypedDict, Union
 
 from numpy import dtype, ndarray, uint32
 from numpy.random.bit_generator import BitGenerator, SeedSequence
 from numpy.typing import _ArrayLikeInt_co
 
-_MT19937State = Dict[str, Union[str, Dict[str, Union[int, ndarray[Any, dtype[uint32]]]]]]
+class _MT19937Internal(TypedDict):
+    key: ndarray[Any, dtype[uint32]]
+    pos: int
+
+class _MT19937State(TypedDict):
+    bit_generator: str
+    state: _MT19937Internal
 
 class MT19937(BitGenerator):
     def __init__(self, seed: Union[None, _ArrayLikeInt_co, SeedSequence] = ...) -> None: ...
diff --git a/numpy/random/_pcg64.pyi b/numpy/random/_pcg64.pyi
index a559604da..c3f9d1769 100644
--- a/numpy/random/_pcg64.pyi
+++ b/numpy/random/_pcg64.pyi
@@ -1,12 +1,17 @@
-from typing import Dict, Union
+from typing import Dict, TypedDict, Union
 
 from numpy.random.bit_generator import BitGenerator, SeedSequence
 from numpy.typing import _ArrayLikeInt_co
 
-_PCG64State = Dict[
-    str,
-    Union[str, int, Dict[str, int]],
-]
+class _PCG64Internal(TypedDict):
+    state: int
+    inc: int
+
+class _PCG64State(TypedDict):
+    bit_generator: str
+    state: _PCG64Internal
+    has_uint32: int
+    uinteger: int
 
 class PCG64(BitGenerator):
     def __init__(self, seed: Union[None, _ArrayLikeInt_co, SeedSequence] = ...) -> None: ...
diff --git a/numpy/random/_philox.pyi b/numpy/random/_philox.pyi
index f622862d6..9bfff87b4 100644
--- a/numpy/random/_philox.pyi
+++ b/numpy/random/_philox.pyi
@@ -1,18 +1,20 @@
-from typing import Any, Dict, Union
+from typing import Any, Dict, TypedDict, Union
 
 from numpy import dtype, ndarray, uint64
 from numpy.random.bit_generator import BitGenerator, SeedSequence
 from numpy.typing import _ArrayLikeInt_co
 
-_PhiloxState = Dict[
-    str,
-    Union[
-        str,
-        int,
-        ndarray[Any, dtype[uint64]],
-        Dict[str, ndarray[Any, dtype[uint64]]],
-    ],
-]
+class _PhiloxInternal(TypedDict):
+    counter: ndarray[Any, dtype[uint64]]
+    key: ndarray[Any, dtype[uint64]]
+
+class _PhiloxState(TypedDict):
+    bit_generator: str
+    state: _PhiloxInternal
+    buffer: ndarray[Any, dtype[uint64]]
+    buffer_pos: int
+    has_uint32: int
+    uinteger: int
 
 class Philox(BitGenerator):
     def __init__(
diff --git a/numpy/random/_sfc64.pyi b/numpy/random/_sfc64.pyi
index 9429011f8..2ccac852e 100644
--- a/numpy/random/_sfc64.pyi
+++ b/numpy/random/_sfc64.pyi
@@ -1,4 +1,4 @@
-from typing import Any, Dict, Union
+from typing import Any, Dict, TypedDict, Union
 
 from numpy import dtype as dtype
 from numpy import ndarray as ndarray
@@ -6,10 +6,14 @@ from numpy import uint64
 from numpy.random.bit_generator import BitGenerator, SeedSequence
 from numpy.typing import _ArrayLikeInt_co
 
-_SFC64State = Dict[
-    str,
-    Union[str, int, Dict[str, ndarray[Any, dtype[uint64]]]],
-]
+class _SFC64Internal(TypedDict):
+    state: ndarray[Any, dtype[uint64]]
+
+class _SFC64State(TypedDict):
+    bit_generator: str
+    state: _SFC64Internal
+    has_uint32: int
+    uinteger: int
 
 class SFC64(BitGenerator):
     def __init__(self, seed: Union[None, _ArrayLikeInt_co, SeedSequence] = ...) -> None: ...
diff --git a/numpy/random/bit_generator.pyi b/numpy/random/bit_generator.pyi
index 6e86cc8e4..f64d42b78 100644
--- a/numpy/random/bit_generator.pyi
+++ b/numpy/random/bit_generator.pyi
@@ -6,6 +6,7 @@ from typing import (
     Callable,
     Dict,
     List,
+    Mapping,
     NamedTuple,
     Optional,
     Sequence,
@@ -99,9 +100,9 @@ class BitGenerator:
         self,
     ) -> Tuple[Callable[[str], BitGenerator], Tuple[str], Tuple[Dict[str, Any]]]: ...
     @property
-    def state(self) -> Dict[str, Any]: ...
+    def state(self) -> Mapping[str, Any]: ...
     @state.setter
-    def state(self, value: Dict[str, Any]) -> None: ...
+    def state(self, value: Mapping[str, Any]) -> None: ...
     @overload
     def random_raw(self, size: None = ..., output: Literal[True] = ...) -> int: ...  # type: ignore[misc]
     @overload