Add RandomState support for df sampling (#526)

* Add RandomState support for df sampling

* Add typing and styling fixes

* Add support for groupby and series

* Remove arbitrary space

Author: intermezzio

pandas-stubs/_typing.pyi

@@ -373,4 +373,13 @@ ValidationOptions: TypeAlias = Literal[
     "many_to_many",
     "m:m",
 ]
+
+RandomState: TypeAlias = Union[
+    int,
+    ArrayLike,
+    np.random.Generator,
+    np.random.BitGenerator,
+    np.random.RandomState,
+]
+
 __all__ = ["npt", "type_t"]

pandas-stubs/core/frame.pyi

@@ -94,6 +94,7 @@ from pandas._typing import (
     NaPosition,
     ParquetEngine,
     QuantileInterpolation,
+    RandomState,
     ReadBuffer,
     Renamer,
     ReplaceMethod,
@@ -1916,7 +1917,7 @@ class DataFrame(NDFrame, OpsMixin):
         frac: float | None = ...,
         replace: _bool = ...,
         weights: _str | ListLike | None = ...,
-        random_state: int | None = ...,
+        random_state: RandomState | None = ...,
         axis: SeriesAxisType | None = ...,
         ignore_index: _bool = ...,
     ) -> DataFrame: ...

pandas-stubs/core/groupby/generic.pyi

@@ -33,6 +33,7 @@ from pandas._typing import (
     AxisType,
     Level,
     ListLike,
+    RandomState,
     Scalar,
 )
 
@@ -302,7 +303,7 @@ class DataFrameGroupBy(GroupBy):
         frac: float | None = ...,
         replace: bool = ...,
         weights: ListLike | None = ...,
-        random_state: int | None = ...,
+        random_state: RandomState | None = ...,
     ) -> DataFrame: ...
     def sem(self, ddof: int = ..., numeric_only: bool = ...) -> DataFrame: ...
     def shift(

pandas-stubs/core/series.pyi

@@ -100,6 +100,7 @@ from pandas._typing import (
     MaskType,
     NaPosition,
     QuantileInterpolation,
+    RandomState,
     Renamer,
     ReplaceMethod,
     Scalar,
@@ -1021,7 +1022,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         frac: float | None = ...,
         replace: _bool = ...,
         weights: _str | _ListLike | np.ndarray | None = ...,
-        random_state: int | None = ...,
+        random_state: RandomState | None = ...,
         axis: SeriesAxisType | None = ...,
         ignore_index: _bool = ...,
     ) -> Series[S1]: ...

tests/test_frame.py

@@ -231,6 +231,10 @@ def test_types_sample() -> None:
     # GH 67
     check(assert_type(df.sample(frac=0.5), pd.DataFrame), pd.DataFrame)
     check(assert_type(df.sample(n=1), pd.DataFrame), pd.DataFrame)
+    check(
+        assert_type(df.sample(n=1, random_state=np.random.default_rng()), pd.DataFrame),
+        pd.DataFrame,
+    )
 
 
 def test_types_nlargest_nsmallest() -> None: