fix: #1008 columns in DataFrame.drop (#1131)

* fix: #1008 columns in DataFrame.drop

* fix(comment): #1008
- #1131 (review)

Author: cmp0xff

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 = ...,

tests/test_frame.py

@@ -344,9 +344,11 @@ 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)
-    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