Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ...,
Expand Down
2 changes: 1 addition & 1 deletion pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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]: ...
Expand Down
11 changes: 11 additions & 0 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
7 changes: 7 additions & 0 deletions tests/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down