diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 203190707..6a946e267 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -577,7 +577,9 @@ class DataFrame(NDFrame, OpsMixin): @overload def set_index( self, - keys: Union[Label, Sequence], + keys: Union[ + Label, Series, Index, np.ndarray, Iterator[Hashable], List[Hashable] + ], drop: _bool = ..., append: _bool = ..., verify_integrity: _bool = ..., @@ -587,7 +589,9 @@ class DataFrame(NDFrame, OpsMixin): @overload def set_index( self, - keys: Union[Label, Sequence], + keys: Union[ + Label, Series, Index, np.ndarray, Iterator[Hashable], List[Hashable] + ], drop: _bool = ..., append: _bool = ..., verify_integrity: _bool = ..., @@ -597,7 +601,9 @@ class DataFrame(NDFrame, OpsMixin): @overload def set_index( self, - keys: Union[Label, Sequence], + keys: Union[ + Label, Series, Index, np.ndarray, Iterator[Hashable], List[Hashable] + ], drop: _bool = ..., append: _bool = ..., *, @@ -606,7 +612,9 @@ class DataFrame(NDFrame, OpsMixin): @overload def set_index( self, - keys: Union[Label, Sequence], + keys: Union[ + Label, Series, Index, np.ndarray, Iterator[Hashable], List[Hashable] + ], drop: _bool = ..., append: _bool = ..., inplace: Optional[_bool] = ..., diff --git a/tests/test_frame.py b/tests/test_frame.py index b76df3b4e..6c799c20e 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -268,6 +268,8 @@ def test_types_set_index() -> None: res4: pd.DataFrame = df.set_index("col1", verify_integrity=True) res5: pd.DataFrame = df.set_index(["col1", "col2"]) res6: None = df.set_index("col1", inplace=True) + # GH 140 + res7: pd.DataFrame = df.set_index(pd.Index(["w", "x", "y", "z"])) def test_types_query() -> None: