TYP: RandomState alias (pandas-dev#42171)

Author: mzeitlin11

pandas/_typing.py

@@ -122,6 +122,7 @@
 JSONSerializable = Optional[Union[PythonScalar, List, Dict]]
 Frequency = Union[str, "DateOffset"]
 Axes = Collection[Any]
+RandomState = Union[int, ArrayLike, np.random.Generator, np.random.RandomState]
 
 # dtypes
 NpDtype = Union[str, np.dtype]

pandas/core/common.py

@@ -30,6 +30,7 @@
 from pandas._typing import (
     AnyArrayLike,
     NpDtype,
+    RandomState,
     Scalar,
     T,
 )
@@ -388,13 +389,13 @@ def standardize_mapping(into):
     return into
 
 
-def random_state(state=None):
+def random_state(state: RandomState | None = None) -> np.random.RandomState:
     """
     Helper function for processing random_state arguments.
 
     Parameters
     ----------
-    state : int, array-like, BitGenerator (NumPy>=1.17), np.random.RandomState, None.
+    state : int, array-like, BitGenerator, np.random.RandomState, None.
         If receives an int, array-like, or BitGenerator, passes to
         np.random.RandomState() as seed.
         If receives an np.random.RandomState object, just returns object.
@@ -403,8 +404,8 @@ def random_state(state=None):
 
         .. versionchanged:: 1.1.0
 
-            array-like and BitGenerator (for NumPy>=1.18) object now passed to
-            np.random.RandomState() as seed
+            array-like and BitGenerator object now passed to np.random.RandomState()
+            as seed
 
         Default None.
 

pandas/core/generic.py

@@ -47,6 +47,7 @@
     Level,
     Manager,
     NpDtype,
+    RandomState,
     Renamer,
     StorageOptions,
     T,
@@ -5146,7 +5147,7 @@ def sample(
         frac=None,
         replace=False,
         weights=None,
-        random_state=None,
+        random_state: RandomState | None = None,
         axis=None,
     ) -> FrameOrSeries:
         """
@@ -5177,8 +5178,7 @@ def sample(
             Missing values in the weights column will be treated as zero.
             Infinite values not allowed.
         random_state : int, array-like, BitGenerator, np.random.RandomState, optional
-            If int, array-like, or BitGenerator (NumPy>=1.17), seed for
-            random number generator
+            If int, array-like, or BitGenerator, seed for random number generator.
             If np.random.RandomState, use as numpy RandomState object.
 
             .. versionchanged:: 1.1.0

pandas/core/groupby/groupby.py

@@ -47,6 +47,7 @@ class providing the base-class of operations.
     FrameOrSeries,
     FrameOrSeriesUnion,
     IndexLabel,
+    RandomState,
     Scalar,
     T,
     final,
@@ -3189,7 +3190,7 @@ def sample(
         frac: float | None = None,
         replace: bool = False,
         weights: Sequence | Series | None = None,
-        random_state=None,
+        random_state: RandomState | None = None,
     ):
         """
         Return a random sample of items from each group.
@@ -3216,8 +3217,7 @@ def sample(
             Values must be non-negative with at least one positive element
             within each group.
         random_state : int, array-like, BitGenerator, np.random.RandomState, optional
-            If int, array-like, or BitGenerator (NumPy>=1.17), seed for
-            random number generator
+            If int, array-like, or BitGenerator, seed for random number generator.
             If np.random.RandomState, use as numpy RandomState object.
 
         Returns