added datatypes for dropna subset as per pandas-dev#434

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

@@ -2332,3 +2332,14 @@ 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:
+    data = {'col1': [1, 3, 4], 'col2': [2, 3, 5], 'col2': [2, 4, 4]}
+    df = pd.DataFrame(data)
+    check(
+        assert_type(
+            df.dropna(subset=df.columns.drop('col1')), pd.DataFrame
+        ),
+        pd.DataFrame,
+    )
+