fix type signatures of Dataframe.squeeze and Series.squeeze

Author: yangdanny97

pandas-stubs/core/frame.pyi

@@ -2191,7 +2191,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
         numeric_only: _bool = ...,
         **kwargs: Any,
     ) -> Series: ...
-    def squeeze(self, axis: Axis | None = ...): ...
+    def squeeze(self, axis: Axis | None = ...) -> DataFrame | Series | Scalar: ...
     def std(
         self,
         axis: Axis = ...,

pandas-stubs/core/series.pyi

@@ -1216,7 +1216,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[S1]: ...
     def droplevel(self, level: Level | list[Level], axis: AxisIndex = ...) -> Self: ...
     def pop(self, item: Hashable) -> S1: ...
-    def squeeze(self, axis: AxisIndex | None = ...) -> Scalar: ...
+    def squeeze(self) -> Series[S1] | Scalar: ...
     def __abs__(self) -> Series[S1]: ...
     def add_prefix(self, prefix: _str, axis: AxisIndex | None = ...) -> Series[S1]: ...
     def add_suffix(self, suffix: _str, axis: AxisIndex | None = ...) -> Series[S1]: ...

tests/test_frame.py

@@ -3231,6 +3231,17 @@ def test_resample() -> None:
     check(assert_type(df.resample("2min").ohlc(), pd.DataFrame), pd.DataFrame)
 
 
+def test_squeeze() -> None:
+    df1 = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
+    check(df1.squeeze(), pd.DataFrame)
+    df2 = pd.DataFrame({"a": [1, 2]})
+    check(df2.squeeze(), pd.Series)
+    df3 = pd.DataFrame({"a": [1], "b": [2]})
+    check(df3.squeeze(), pd.Series, np.integer)
+    df4 = pd.DataFrame({"a": [1]})
+    check(df4.squeeze(), np.integer)
+
+
 def test_loc_set() -> None:
     df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
     df.loc["a"] = [3, 4]

tests/test_series.py

@@ -1730,6 +1730,13 @@ def test_resample() -> None:
     check(assert_type(s.resample("2min").ohlc(), pd.DataFrame), pd.DataFrame)
 
 
+def test_squeeze() -> None:
+    s1 = pd.Series([1, 2, 3])
+    check(s1.squeeze(), pd.Series, np.integer)
+    s2 = pd.Series([1])
+    check(s2.squeeze(), np.integer)
+
+
 def test_to_xarray():
     s = pd.Series([1, 2])
     check(assert_type(s.to_xarray(), xr.DataArray), xr.DataArray)