From e6764019a6321154bbdc628e4dd58c4595eb2369 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Wed, 26 Feb 2025 23:03:30 +0100 Subject: [PATCH 1/2] fix: #1008 columns in DataFrame.drop --- pandas-stubs/core/generic.pyi | 8 ++++---- tests/test_frame.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi index 5aeae8c6c..0c432b86b 100644 --- a/pandas-stubs/core/generic.pyi +++ b/pandas-stubs/core/generic.pyi @@ -289,7 +289,7 @@ class NDFrame(indexing.IndexingMixin): *, axis: Axis = ..., index: Hashable | Sequence[Hashable] | Index[Any] = ..., - columns: Hashable | Sequence[Hashable] | Index[Any], + columns: Hashable | Iterable[Hashable], level: Level | None = ..., inplace: Literal[True], errors: IgnoreRaise = ..., @@ -301,7 +301,7 @@ class NDFrame(indexing.IndexingMixin): *, axis: Axis = ..., index: Hashable | Sequence[Hashable] | Index[Any], - columns: Hashable | Sequence[Hashable] | Index[Any] = ..., + columns: Hashable | Iterable[Hashable] = ..., level: Level | None = ..., inplace: Literal[True], errors: IgnoreRaise = ..., @@ -325,7 +325,7 @@ class NDFrame(indexing.IndexingMixin): *, axis: Axis = ..., index: Hashable | Sequence[Hashable] | Index[Any] = ..., - columns: Hashable | Sequence[Hashable] | Index[Any], + columns: Hashable | Iterable[Hashable], level: Level | None = ..., inplace: Literal[False] = ..., errors: IgnoreRaise = ..., @@ -337,7 +337,7 @@ class NDFrame(indexing.IndexingMixin): *, axis: Axis = ..., index: Hashable | Sequence[Hashable] | Index[Any], - columns: Hashable | Sequence[Hashable] | Index[Any] = ..., + columns: Hashable | Iterable[Hashable] = ..., level: Level | None = ..., inplace: Literal[False] = ..., errors: IgnoreRaise = ..., diff --git a/tests/test_frame.py b/tests/test_frame.py index 7696008e3..0fa2d9089 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -344,9 +344,9 @@ def test_types_drop() -> None: df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) check(assert_type(df.drop("col1", axis=1), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop(columns=["col1"]), pd.DataFrame), pd.DataFrame) + check(assert_type(df.drop(columns=iter(["col1"])), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop([0]), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop(index=[0]), pd.DataFrame), pd.DataFrame) - check(assert_type(df.drop(columns=["col1"]), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop(index=1), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop(labels=0), pd.DataFrame), pd.DataFrame) assert assert_type(df.drop([0, 0], inplace=True), None) is None From 2697a018e8f5f6ccad3b32cf2ec72f4e4c525b81 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Thu, 27 Feb 2025 08:24:51 +0100 Subject: [PATCH 2/2] fix(comment): #1008 - https://github.com/pandas-dev/pandas-stubs/pull/1131#pullrequestreview-2646108815 --- tests/test_frame.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_frame.py b/tests/test_frame.py index 0fa2d9089..e506528e3 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -344,6 +344,8 @@ def test_types_drop() -> None: df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) check(assert_type(df.drop("col1", axis=1), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop(columns=["col1"]), pd.DataFrame), pd.DataFrame) + check(assert_type(df.drop(columns=pd.Index(["col1"])), pd.DataFrame), pd.DataFrame) + check(assert_type(df.drop(columns={"col1"}), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop(columns=iter(["col1"])), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop([0]), pd.DataFrame), pd.DataFrame) check(assert_type(df.drop(index=[0]), pd.DataFrame), pd.DataFrame)