ENH: updating subset types in drop_duplicates/duplicated (#986)

* updating subset types

* adding pytest

* updating based on feedback

* updating based on feedback

* updating based on feedback

* correcting conditional

Co-authored-by: Irv Lustig <[email protected]>

* complying with black formatting

---------

Co-authored-by: Irv Lustig <[email protected]>

Author: kirill-bash

pandas-stubs/core/frame.pyi

@@ -901,15 +901,15 @@ class DataFrame(NDFrame, OpsMixin):
     ) -> DataFrame | None: ...
     def drop_duplicates(
         self,
-        subset=...,
+        subset: Hashable | Iterable[Hashable] | None = ...,
         *,
         keep: NaPosition | _bool = ...,
         inplace: _bool = ...,
         ignore_index: _bool = ...,
     ) -> DataFrame: ...
     def duplicated(
         self,
-        subset: Hashable | Sequence[Hashable] | None = ...,
+        subset: Hashable | Iterable[Hashable] | None = ...,
         keep: NaPosition | _bool = ...,
     ) -> Series: ...
     @overload

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)