fix: DataFrame.sample argument 'axis' can be any Axis (#1252)

nithinmanne · web-flow · commit f745945dfcd0 · 2025-06-16T09:43:18.000-04:00
diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
@@ -2272,7 +2272,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
         replace: _bool = ...,
         weights: _str | ListLike | None = ...,
         random_state: RandomState | None = ...,
-        axis: AxisIndex | None = ...,
+        axis: Axis | None = ...,
         ignore_index: _bool = ...,
     ) -> Self: ...
     def sem(
diff --git a/tests/test_frame.py b/tests/test_frame.py
@@ -370,6 +370,7 @@ def test_types_sample() -> None:
         assert_type(df.sample(n=1, random_state=np.random.default_rng()), pd.DataFrame),
         pd.DataFrame,
     )
+    check(assert_type(df.sample(n=1, axis=1), pd.DataFrame), pd.DataFrame)
 
 
 def test_types_nlargest_nsmallest() -> None:

Original file line number	Diff line number	Diff line change
`@@ -370,6 +370,7 @@ def test_types_sample() -> None:`
`370`	`370`	`assert_type(df.sample(n=1, random_state=np.random.default_rng()), pd.DataFrame),`
`371`	`371`	`pd.DataFrame,`
`372`	`372`	`)`
	`373`	`+ check(assert_type(df.sample(n=1, axis=1), pd.DataFrame), pd.DataFrame)`
`373`	`374`
`374`	`375`
`375`	`376`	`def test_types_nlargest_nsmallest() -> None:`