ENH: Add typing for MT19937 and Philox

Add typing for remaining core bit generators
Clean up remaining files

6 files changed, 92 insertions, 61 deletions

diff --git a/numpy/random/__init__.pyi b/numpy/random/__init__.pyi
index d8b7a2827..6507a06a4 100644
--- a/numpy/random/__init__.pyi
+++ b/numpy/random/__init__.pyi
@@ -1,8 +1,11 @@
 from typing import Any, List
-from numpy.random.bit_generator import (
-    BitGenerator as BitGenerator,
-    SeedSequence as SeedSequence
-)
+
+from numpy.random._mt19937 import MT19937 as MT19937
+from numpy.random._pcg64 import PCG64 as PCG64
+from numpy.random._philox import Philox as Philox
+from numpy.random._sfc64 import SFC64 as SFC64
+from numpy.random.bit_generator import BitGenerator as BitGenerator
+from numpy.random.bit_generator import SeedSequence as SeedSequence
 
 __all__: List[str]
 
@@ -58,8 +61,4 @@ weibull: Any
 zipf: Any
 Generator: Any
 RandomState: Any
-MT19937: Any
-Philox: Any
-PCG64: Any
-SFC64: Any
 default_rng: Any
diff --git a/numpy/random/_mt19937.pyi b/numpy/random/_mt19937.pyi
new file mode 100644
index 000000000..91eccda38
--- /dev/null
+++ b/numpy/random/_mt19937.pyi
@@ -0,0 +1,16 @@
+from typing import Any, Dict, 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 MT19937(BitGenerator):
+    def __init__(self, seed: Union[None, _ArrayLikeInt_co, SeedSequence] = ...) -> None: ...
+    def _legacy_seeding(self, seed: _ArrayLikeInt_co) -> None: ...
+    def jumped(self, jumps: int = ...) -> MT19937: ...
+    @property
+    def state(self) -> _MT19937State: ...
+    @state.setter
+    def state(self, value: _MT19937State) -> None: ...
diff --git a/numpy/random/_pcg64.pyi b/numpy/random/_pcg64.pyi
index 591459374..a559604da 100644
--- a/numpy/random/_pcg64.pyi
+++ b/numpy/random/_pcg64.pyi
@@ -1,27 +1,23 @@
-from typing import Dict, Literal, Union
+from typing import Dict, Union
 
-from numpy.random.bit_generator import BitGenerator as BitGenerator
-from numpy.random.bit_generator import SeedSequence as SeedSequence
+from numpy.random.bit_generator import BitGenerator, SeedSequence
 from numpy.typing import _ArrayLikeInt_co
 
+_PCG64State = Dict[
+    str,
+    Union[str, int, Dict[str, int]],
+]
+
 class PCG64(BitGenerator):
-    def __init__(
-        self, seed: Union[None, int, _ArrayLikeInt_co, SeedSequence] = ...
-    ) -> None: ...
-    def jumped(self, jumps=1) -> PCG64: ...
+    def __init__(self, seed: Union[None, _ArrayLikeInt_co, SeedSequence] = ...) -> None: ...
+    def jumped(self, jumps: int = ...) -> PCG64: ...
     @property
     def state(
         self,
-    ) -> Dict[
-        Literal["bit_generator", "state", "has_uint32", "uinteger"],
-        Union[str, int, Dict[str, int]],
-    ]: ...
+    ) -> _PCG64State: ...
     @state.setter
     def state(
         self,
-        value: Dict[
-            Literal["bit_generator", "state", "has_uint32", "uinteger"],
-            Union[str, int, Dict[str, int]],
-        ],
+        value: _PCG64State,
     ) -> None: ...
     def advance(self, delta: int) -> PCG64: ...
diff --git a/numpy/random/_philox.pyi b/numpy/random/_philox.pyi
new file mode 100644
index 000000000..f622862d6
--- /dev/null
+++ b/numpy/random/_philox.pyi
@@ -0,0 +1,34 @@
+from typing import Any, Dict, 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 Philox(BitGenerator):
+    def __init__(
+        self,
+        seed: Union[None, _ArrayLikeInt_co, SeedSequence] = ...,
+        counter: Union[None, _ArrayLikeInt_co] = ...,
+        key: Union[None, _ArrayLikeInt_co] = ...,
+    ) -> None: ...
+    @property
+    def state(
+        self,
+    ) -> _PhiloxState: ...
+    @state.setter
+    def state(
+        self,
+        value: _PhiloxState,
+    ) -> None: ...
+    def jumped(self, jumps: int = ...) -> Philox: ...
+    def advance(self, delta: int) -> Philox: ...
diff --git a/numpy/random/_sfc64.pyi b/numpy/random/_sfc64.pyi
index 43190657c..9429011f8 100644
--- a/numpy/random/_sfc64.pyi
+++ b/numpy/random/_sfc64.pyi
@@ -1,27 +1,24 @@
-from typing import Any, Dict, Literal, Union
+from typing import Any, Dict, Union
 
 from numpy import dtype as dtype
 from numpy import ndarray as ndarray
 from numpy import uint64
-from numpy.random.bit_generator import BitGenerator as BitGenerator
-from numpy.random.bit_generator import SeedSequence as SeedSequence
+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 SFC64(BitGenerator):
-    def __init__(
-        self, seed: Union[None, int, _ArrayLikeInt_co, SeedSequence] = ...
-    ) -> None: ...
+    def __init__(self, seed: Union[None, _ArrayLikeInt_co, SeedSequence] = ...) -> None: ...
+    @property
     def state(
         self,
-    ) -> Dict[
-        Literal["bit_generator", "state", "has_uint32", "uinteger"],
-        Union[str, int, Dict[str, ndarray[Any, dtype[uint64]]]],
-    ]: ...
+    ) -> _SFC64State: ...
     @state.setter
     def state(
         self,
-        value: Dict[
-            Literal["bit_generator", "state", "has_uint32", "uinteger"],
-            Union[str, int, Dict[str, ndarray[Any, dtype[uint64]]]],
-        ],
+        value: _SFC64State,
     ) -> None: ...
diff --git a/numpy/random/bit_generator.pyi b/numpy/random/bit_generator.pyi
index 90c7cb08f..96ffbec13 100644
--- a/numpy/random/bit_generator.pyi
+++ b/numpy/random/bit_generator.pyi
@@ -1,11 +1,11 @@
 import abc
+import sys
 from threading import Lock
 from typing import (
     Any,
     Callable,
     Dict,
     List,
-    Literal,
     NamedTuple,
     Optional,
     Sequence,
@@ -17,13 +17,12 @@ from typing import (
 )
 
 from numpy import dtype, ndarray, uint32, uint64, unsignedinteger
-from numpy.typing import (
-    DTypeLike,
-    _ArrayLikeInt_co,
-    _DTypeLikeUInt,
-    _ShapeLike,
-    _SupportsDType,
-)
+from numpy.typing import DTypeLike, _ArrayLikeInt_co, _DTypeLikeUInt, _ShapeLike, _SupportsDType
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
 
 _T = TypeVar("_T")
 
@@ -71,7 +70,7 @@ class SeedlessSeedSequence(ISpawnableSeedSequence):
 
 class SeedSequence(ISpawnableSeedSequence):
     entropy: Union[None, int, Sequence[int]]
-    spawn_key: Tuple[int, ...]
+    spawn_key: Sequence[int]
     pool_size: int
     n_children_spawned: int
     pool: ndarray[Any, dtype[uint32]]
@@ -88,23 +87,17 @@ class SeedSequence(ISpawnableSeedSequence):
     def state(
         self,
     ) -> Dict[str, Union[None, Sequence[int], int, Tuple[int, ...]]]: ...
-    def generate_state(
-        self, n_words: int, dtype: DTypeLike = ...
-    ) -> ndarray[Any, Any]: ...
+    def generate_state(self, n_words: int, dtype: DTypeLike = ...) -> ndarray[Any, Any]: ...
     def spawn(self, n_children: int) -> List[SeedSequence]: ...
 
 class BitGenerator:
     lock: Lock
-    def __init__(
-        self, seed: Union[None, int, _ArrayLikeInt_co, SeedSequence] = ...
-    ) -> None: ...
+    def __init__(self, seed: Union[None, _ArrayLikeInt_co, SeedSequence] = ...) -> None: ...
     def __getstate__(self) -> Dict[str, Any]: ...
     def __setstate__(self, state: Dict[str, Any]) -> None: ...
     def __reduce__(
         self,
-    ) -> Tuple[
-        Callable[[str], BitGenerator], Tuple[str], Tuple[Dict[str, Any]]
-    ]: ...
+    ) -> Tuple[Callable[[str], BitGenerator], Tuple[str], Tuple[Dict[str, Any]]]: ...
     @property
     def state(self) -> Dict[str, Any]: ...
     @state.setter
@@ -112,13 +105,9 @@ class BitGenerator:
     @overload
     def random_raw(self, size: None = ..., output: Literal[True] = ...) -> int: ...  # type: ignore[misc]
     @overload
-    def random_raw(
-        self, size: _ShapeLike = ..., output: Literal[True] = ...
-    ) -> ndarray[Any, dtype[uint64]]: ...  # type: ignore[misc]
+    def random_raw(self, size: _ShapeLike = ..., output: Literal[True] = ...) -> ndarray[Any, dtype[uint64]]: ...  # type: ignore[misc]
     @overload
-    def random_raw(
-        self, size: Optional[_ShapeLike] = ..., output: Literal[False] = ...
-    ) -> None: ...  # type: ignore[misc]
+    def random_raw(self, size: Optional[_ShapeLike] = ..., output: Literal[False] = ...) -> None: ...  # type: ignore[misc]
     def _benchmark(self, cnt: int, method: str = ...) -> None: ...
     @property
     def ctypes(self) -> _Interface: ...