diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 0e151c4a6..5a7438feb 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -901,7 +901,7 @@ class DataFrame(NDFrame, OpsMixin): ) -> DataFrame | None: ... def drop_duplicates( self, - subset=..., + subset: Hashable | Iterable[Hashable] | None = ..., *, keep: NaPosition | _bool = ..., inplace: _bool = ..., @@ -909,7 +909,7 @@ class DataFrame(NDFrame, OpsMixin): ) -> DataFrame: ... def duplicated( self, - subset: Hashable | Sequence[Hashable] | None = ..., + subset: Hashable | Iterable[Hashable] | None = ..., keep: NaPosition | _bool = ..., ) -> Series: ... @overload diff --git a/tests/test_frame.py b/tests/test_frame.py index 33aa1616e..43082143a 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -368,6 +368,28 @@ def test_types_dropna() -> None: res3: None = df.dropna(axis=0, how="all", subset=["col1"], inplace=True) +def test_types_drop_duplicates() -> None: + # GH#59237 + df = pd.DataFrame( + { + "AAA": ["foo", "bar", "foo", "bar", "foo", "bar", "bar", "foo"], + "B": ["one", "one", "two", "two", "two", "two", "one", "two"], + "C": [1, 1, 2, 2, 2, 2, 1, 2], + "D": range(8), + } + ) + + check(assert_type(df.drop_duplicates(["AAA"]), pd.DataFrame), pd.DataFrame) + check(assert_type(df.drop_duplicates(("AAA",)), pd.DataFrame), pd.DataFrame) + check(assert_type(df.drop_duplicates("AAA"), pd.DataFrame), pd.DataFrame) + + if not PD_LTE_22: + check(assert_type(df.drop_duplicates({"AAA"}), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.drop_duplicates({"AAA": None}), pd.DataFrame), pd.DataFrame + ) + + def test_types_fillna() -> None: df = pd.DataFrame(data={"col1": [np.nan, np.nan], "col2": [3, np.nan]}) res: pd.DataFrame = df.fillna(0)