ENH: Add typing for Generator

4 files changed, 253 insertions, 8 deletions

diff --git a/numpy/random/__init__.pyi b/numpy/random/__init__.pyi
index 6507a06a4..b99c002ae 100644
--- a/numpy/random/__init__.pyi
+++ b/numpy/random/__init__.pyi
@@ -1,5 +1,7 @@
 from typing import Any, List
 
+from numpy.random._generator import Generator as Generator
+from numpy.random._generator import default_rng as default_rng
 from numpy.random._mt19937 import MT19937 as MT19937
 from numpy.random._pcg64 import PCG64 as PCG64
 from numpy.random._philox import Philox as Philox
@@ -59,6 +61,4 @@ vonmises: Any
 wald: Any
 weibull: Any
 zipf: Any
-Generator: Any
 RandomState: Any
-default_rng: Any
diff --git a/numpy/random/_generator.pyi b/numpy/random/_generator.pyi
new file mode 100644
index 000000000..22860b5ea
--- /dev/null
+++ b/numpy/random/_generator.pyi
@@ -0,0 +1,245 @@
+import sys
+from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union, overload
+
+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
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+class Generator:
+    # COMPLETE
+    _bit_generator: BitGenerator
+    _poisson_lam_max: float
+    def __init__(self, bit_generator: BitGenerator) -> None: ...
+    def __repr__(self) -> str: ...
+    def __str__(self) -> str: ...
+    # Pickling support:
+    def __getstate__(self) -> Dict[str, Any]: ...
+    def __setstate__(self, state: Dict[str, Any]): ...
+    def __reduce__(self) -> Tuple[Callable[[str], BitGenerator], Tuple[str], Dict[str, Any]]: ...
+    @property
+    def bit_generator(self) -> BitGenerator: ...
+    def bytes(self, length: int) -> str: ...
+    # TODO: Needs overloading
+    def standard_cauchy(
+        self, size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    # TODO: Needs overloading and specific dtypes
+    def standard_exponential(
+        self,
+        size: Optional[Union[_ShapeLike]] = ...,
+        dtype: DTypeLike = ...,
+        method: Literal["zig", "inv"] = ...,
+        out=Union[None, ndarray[Any, dtype[float32]], ndarray[Any, dtype[float64]]],
+    ) -> Union[float, ndarray[Any, Any]]: ...
+    # TODO: Needs typing
+    def random(
+        self,
+        size: Optional[_ShapeLike] = ...,
+        dtype: DTypeLike = ...,
+        out: Union[None, ndarray[Any, dtype[float32]], ndarray[Any, dtype[float64]]] = ...,
+    ): ...
+    def beta(
+        self, a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ): ...
+    def exponential(self, scale: _ArrayLikeFloat_co = ..., size: Union[None, _ShapeLike] = ...): ...
+    def integers(
+        self,
+        low: _ArrayLikeInt_co,
+        high: Optional[_ArrayLikeInt_co] = ...,
+        size: Union[None, _ShapeLike] = ...,
+        dtype=...,
+        endpoint: bool = ...,
+    ) -> ndarray[Any, dtype[integer]]: ...
+    # TODO: Use a TypeVar _T here to get away from Any output?  Should be int->ndarray[Any,dtype[int64]], ArrayLike[_T] -> Union[_T, ndarray[Any,Any]]
+    def choice(
+        self,
+        a: ArrayLike,
+        size: Union[None, _ShapeLike] = ...,
+        replace: bool = ...,
+        p=...,
+        axis=...,
+        shuffle: bool = ...,
+    ) -> Any: ...
+    def uniform(
+        self,
+        low: _ArrayLikeFloat_co = ...,
+        high: _ArrayLikeFloat_co = ...,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    @overload
+    def standard_normal(
+        self,
+        size: None = ...,
+        dtype: DTypeLike = ...,
+        out: None = ...,
+    ) -> float: ...
+    # TODO: How to literal dtype?
+    @overload
+    def standard_normal(
+        self,
+        size: _ShapeLike = ...,
+        dtype: DTypeLike = ...,
+        out: Union[None, ndarray[Any, dtype[float32]], ndarray[Any, dtype[float64]]] = ...,
+    ) -> Union[ndarray[Any, dtype[float32]], ndarray[Any, dtype[float64]]]: ...
+    def normal(
+        self,
+        loc: _ArrayLikeFloat_co = ...,
+        scale: _ArrayLikeFloat_co = ...,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def standard_gamma(
+        self,
+        shape,
+        size: Union[None, _ShapeLike] = ...,
+        dtype=...,
+        out: Union[None, ndarray[Any, dtype[float32]], ndarray[Any, dtype[float64]]] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def gamma(
+        self,
+        shape: _ArrayLikeFloat_co,
+        scale: _ArrayLikeFloat_co = ...,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def f(
+        self,
+        dfnum: _ArrayLikeFloat_co,
+        dfden: _ArrayLikeFloat_co,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def noncentral_f(
+        self,
+        dfnum: _ArrayLikeFloat_co,
+        dfden: _ArrayLikeFloat_co,
+        nonc: _ArrayLikeFloat_co,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def chisquare(
+        self, df: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def noncentral_chisquare(
+        self, df: _ArrayLikeFloat_co, nonc: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def standard_t(
+        self, df: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def vonmises(
+        self, mu: _ArrayLikeFloat_co, kappa: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def pareto(
+        self, a: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def weibull(
+        self, a: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def power(
+        self, a: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def laplace(
+        self,
+        loc: _ArrayLikeFloat_co = ...,
+        scale: _ArrayLikeFloat_co = ...,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def gumbel(
+        self,
+        loc: _ArrayLikeFloat_co = ...,
+        scale: _ArrayLikeFloat_co = ...,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def logistic(
+        self,
+        loc: _ArrayLikeFloat_co = ...,
+        scale: _ArrayLikeFloat_co = ...,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def lognormal(
+        self,
+        mean: _ArrayLikeFloat_co = ...,
+        sigma: _ArrayLikeFloat_co = ...,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def rayleigh(
+        self, scale: _ArrayLikeFloat_co = ..., size: Union[None, _ShapeLike] = ...
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def wald(
+        self,
+        mean: _ArrayLikeFloat_co,
+        scale: _ArrayLikeFloat_co,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    def triangular(
+        self,
+        left: _ArrayLikeFloat_co,
+        mode: _ArrayLikeFloat_co,
+        right: _ArrayLikeFloat_co,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[float, ndarray[Any, dtype[float64]]]: ...
+    # Complicated, discrete distributions:
+    def binomial(
+        self, n: _ArrayLikeInt_co, p: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[int, ndarray[Any, dtype[int64]]]: ...
+    def negative_binomial(
+        self, n: _ArrayLikeFloat_co, p: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[int, ndarray[Any, dtype[int64]]]: ...
+    def poisson(
+        self, lam: _ArrayLikeFloat_co = ..., size: Union[None, _ShapeLike] = ...
+    ) -> Union[int, ndarray[Any, dtype[int64]]]: ...
+    def zipf(
+        self, a: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[int, ndarray[Any, dtype[int64]]]: ...
+    def geometric(
+        self, p: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[int, ndarray[Any, dtype[int64]]]: ...
+    def hypergeometric(
+        self,
+        ngood: _ArrayLikeInt_co,
+        nbad: _ArrayLikeInt_co,
+        nsample: _ArrayLikeInt_co,
+        size: Union[None, _ShapeLike] = ...,
+    ) -> Union[int, ndarray[Any, dtype[int64]]]: ...
+    def logseries(
+        self, p: _ArrayLikeFloat_co, size: Union[None, _ShapeLike] = ...
+    ) -> Union[int, ndarray[Any, dtype[int64]]]: ...
+    # Multivariate distributions:
+    # TODO: Really need  1-d array like floating and 2-d array-like floating. Using Sequence[float] ??
+    def multivariate_normal(
+        self,
+        mean: Sequence[float],
+        cov: Sequence[Sequence[float]],
+        size: Union[None, _ShapeLike] = ...,
+        check_valid: Literal["warn", "raise", "ignore"] = ...,
+        tol: float = ...,
+        *,
+        method: Literal["svd", "eigh", "cholesky"] = ...
+    ): ...
+    # TODO: Need 1-d array like floating. Using Sequence[float] ??
+    def multinomial(
+        self, n: _ArrayLikeInt_co, pvals: Sequence[float], size: Union[None, _ShapeLike] = ...
+    ): ...
+    # TODO: Need 1-d array like integers. Using Sequence[int] ??
+    def multivariate_hypergeometric(
+        self,
+        colors: Sequence[int],
+        nsample: int,
+        size: Union[None, _ShapeLike] = ...,
+        method: Literal["marginals", "count"] = ...,
+    ): ...
+    # TODO: Need 1-d array like floating. Using Sequence[float] ??
+    def dirichlet(
+        self, alpha: Sequence[float], size: Union[None, _ShapeLike] = ...
+    ) -> ndarray[Any, dtype[float64]]: ...
+    def permuted(
+        self, x: ArrayLike, *, axis: Optional[int] = ..., out: Optional[ndarray[Any, Any]] = ...
+    ) -> ndarray[Any, Any]: ...
+    def shuffle(self, x: ArrayLike, axis: int = ...) -> Sequence[Any]: ...
+    @overload
+    def permutation(self, x: int, axis: int = ...) -> ndarray[Any, dtype[int64]]: ...
+    @overload
+    def permutation(self, x: ArrayLike, axis: int = ...) -> ndarray[Any, Any]: ...
+
+def default_rng(seed=None) -> Generator: ...
diff --git a/numpy/random/_generator.pyx b/numpy/random/_generator.pyx
index cc7991eb1..a7d98e2ed 100644
--- a/numpy/random/_generator.pyx
+++ b/numpy/random/_generator.pyx
@@ -579,13 +579,13 @@ cdef class Generator:
 
         Returns
         -------
-        out : str
+        out : bytes
             String of length `length`.
 
         Examples
         --------
         >>> np.random.default_rng().bytes(10)
-        ' eh\\x85\\x022SZ\\xbf\\xa4' #random
+        b'\xfeC\x9b\x86\x17\xf2\xa1\xafcp' # random
 
         """
         cdef Py_ssize_t n_uint32 = ((length - 1) // 4 + 1)
@@ -996,7 +996,7 @@ cdef class Generator:
         --------
         >>> rng = np.random.default_rng()
         >>> rng.standard_normal()
-        2.1923875335537315 #random
+        2.1923875335537315 # random
 
         >>> s = rng.standard_normal(8000)
         >>> s
@@ -1755,7 +1755,7 @@ cdef class Generator:
         statistic appear?
 
         >>> np.sum(s<t) / float(len(s))
-        0.0090699999999999999  #random
+        0.0090699999999999999  # random
 
         So the p-value is about 0.009, which says the null hypothesis has a
         probability of about 99% of being true.
@@ -3196,7 +3196,7 @@ cdef class Generator:
         How many trials succeeded after a single run?
 
         >>> (z == 1).sum() / 10000.
-        0.34889999999999999 #random
+        0.34889999999999999 # random
 
         """
         return disc(&random_geometric, &self._bitgen, size, self.lock, 1, 0,
diff --git a/numpy/random/bit_generator.pyi b/numpy/random/bit_generator.pyi
index 96ffbec13..6e86cc8e4 100644
--- a/numpy/random/bit_generator.pyi
+++ b/numpy/random/bit_generator.pyi
@@ -101,7 +101,7 @@ class BitGenerator:
     @property
     def state(self) -> Dict[str, Any]: ...
     @state.setter
-    def state(self, value: Dict[str, Any]): ...
+    def state(self, value: Dict[str, Any]) -> None: ...
     @overload
     def random_raw(self, size: None = ..., output: Literal[True] = ...) -> int: ...  # type: ignore[misc]
     @overload