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..e506528e3 100644 --- a/tests/test_frame.py +++ b/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