Skip to content

Commit 0620024

Browse files
committed
added list[Levels] as per #489 and added appropriate tests
added a generic Any for series output created tests for stack wrapping generics in quotes for compatibility style update
1 parent 0b46778 commit 0620024

File tree

2 files changed

+24
-1
lines changed

2 files changed

+24
-1
lines changed

pandas-stubs/core/frame.pyi

+3-1
Original file line numberDiff line numberDiff line change
@@ -997,7 +997,9 @@ class DataFrame(NDFrame, OpsMixin):
997997
margins_name: _str = ...,
998998
observed: _bool = ...,
999999
) -> DataFrame: ...
1000-
def stack(self, level: Level = ..., dropna: _bool = ...) -> DataFrame | Series: ...
1000+
def stack(
1001+
self, level: Level | list[Level] = ..., dropna: _bool = ...
1002+
) -> DataFrame | Series[Any]: ...
10011003
def explode(self, column: _str | tuple, ignore_index: _bool = ...) -> DataFrame: ...
10021004
def unstack(
10031005
self,

tests/test_frame.py

+21
Original file line numberDiff line numberDiff line change
@@ -1798,6 +1798,27 @@ def test_frame_index_numpy() -> None:
17981798
pd.DataFrame([[1.0, 2.0], [3.0, 4.0]], columns=["a", "b"], index=i)
17991799

18001800

1801+
def test_frame_stack() -> None:
1802+
1803+
multicol2 = pd.MultiIndex.from_tuples([("weight", "kg"), ("height", "m")])
1804+
df_multi_level_cols2 = pd.DataFrame(
1805+
[[1.0, 2.0], [3.0, 4.0]], index=["cat", "dog"], columns=multicol2
1806+
)
1807+
1808+
check(
1809+
assert_type(
1810+
df_multi_level_cols2.stack(0), Union[pd.DataFrame, "pd.Series[Any]"]
1811+
),
1812+
pd.DataFrame,
1813+
)
1814+
check(
1815+
assert_type(
1816+
df_multi_level_cols2.stack([0, 1]), Union[pd.DataFrame, "pd.Series[Any]"]
1817+
),
1818+
pd.Series,
1819+
)
1820+
1821+
18011822
def test_frame_reindex() -> None:
18021823
# GH 84
18031824
df = pd.DataFrame({"a": [1, 2, 3]}, index=[0, 1, 2])

0 commit comments

Comments
 (0)