added datatypes for dropna subset as per #434 (#498)

* added datatypes for dropna subset as per #434

* fixed failures

* Formatted file using poe style

Author: shubhankarunhale

pandas-stubs/core/frame.pyi

@@ -789,7 +789,7 @@ class DataFrame(NDFrame, OpsMixin):
         axis: AxisType = ...,
         how: Literal["any", "all"] = ...,
         thresh: int | None = ...,
-        subset: list | None = ...,
+        subset: ListLikeU | Scalar | None = ...,
         inplace: Literal[True],
     ) -> None: ...
     @overload
@@ -799,7 +799,7 @@ class DataFrame(NDFrame, OpsMixin):
         axis: AxisType = ...,
         how: Literal["any", "all"] = ...,
         thresh: int | None = ...,
-        subset: list | None = ...,
+        subset: ListLikeU | Scalar | None = ...,
         inplace: Literal[False] = ...,
     ) -> DataFrame: ...
     @overload
@@ -809,7 +809,7 @@ class DataFrame(NDFrame, OpsMixin):
         axis: AxisType = ...,
         how: Literal["any", "all"] = ...,
         thresh: int | None = ...,
-        subset: list | None = ...,
+        subset: ListLikeU | Scalar | None = ...,
         inplace: _bool | None = ...,
     ) -> DataFrame | None: ...
     def drop_duplicates(

tests/test_frame.py

@@ -2353,3 +2353,13 @@ def test_getsetitem_multiindex() -> None:
     multi_index: tuple[str, str] = ("Year 1", "Q1")
     budget.loc["project A", multi_index] = 4700
     check(assert_type(budget.loc["project A", multi_index], Scalar), int)
+
+
+def test_frame_dropna_subset() -> None:
+    # GH 434
+    data = {"col1": [1, 3, 4], "col2": [2, 3, 5], "col3": [2, 4, 4]}
+    df = pd.DataFrame(data)
+    check(
+        assert_type(df.dropna(subset=df.columns.drop("col1")), pd.DataFrame),
+        pd.DataFrame,
+    )