Skip to content

Commit

Permalink
fix(frame): add arguments for stack method (#926)
Browse files Browse the repository at this point in the history
* fix(frame): add arguments for stack method

* feat: test future_stack only for LTE 2.2

* fix: include both test in PD_LTE_22 condition

---------

Co-authored-by: Quentin FLEURENT NAMBOT <qfleurentnambot@hevaweb.com>
  • Loading branch information
rinzool and Quentin FLEURENT NAMBOT committed May 15, 2024
1 parent a66ea6b commit 8a41508
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 1 deletion.
7 changes: 6 additions & 1 deletion pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -1191,8 +1191,13 @@ class DataFrame(NDFrame, OpsMixin):
margins_name: _str = ...,
observed: _bool = ...,
) -> DataFrame: ...
@overload
def stack(
self, level: Level | list[Level] = ..., dropna: _bool = ..., sort: _bool = ...
) -> DataFrame | Series[Any]: ...
@overload
def stack(
self, level: Level | list[Level] = ..., dropna: _bool = ...
self, level: Level | list[Level] = ..., future_stack: _bool = ...
) -> DataFrame | Series[Any]: ...
def explode(
self, column: Sequence[Hashable], ignore_index: _bool = ...
Expand Down
15 changes: 15 additions & 0 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -2256,6 +2256,21 @@ def test_frame_stack() -> None:
),
pd.Series,
)
if PD_LTE_22:
check(
assert_type(
df_multi_level_cols2.stack(0, future_stack=False),
Union[pd.DataFrame, "pd.Series[Any]"],
),
pd.DataFrame,
)
check(
assert_type(
df_multi_level_cols2.stack(0, dropna=True, sort=True),
Union[pd.DataFrame, "pd.Series[Any]"],
),
pd.DataFrame,
)


def test_frame_reindex() -> None:
Expand Down

0 comments on commit 8a41508

Please sign in to comment.