Add typing for dropna for Index and MultiIndex (#232)

Author: phofl

pandas-stubs/core/indexes/base.pyi

@@ -128,7 +128,7 @@ class Index(IndexOpsMixin, PandasObject):
     def notna(self): ...
     notnull = ...
     def fillna(self, value=..., downcast=...): ...
-    def dropna(self, how: _str = ...): ...
+    def dropna(self, how: Literal["any", "all"] = ...) -> Index: ...
     def unique(self, level=...) -> Index: ...
     def drop_duplicates(
         self, keep: NaPosition | Literal[False] = ...

pandas-stubs/core/indexes/extension.pyi

@@ -1,9 +1,11 @@
+from typing import Literal
+
 from pandas.core.indexes.base import Index
 
 class ExtensionIndex(Index):
     def __getitem__(self, key): ...
     def __iter__(self): ...
-    def dropna(self, how: str = ...): ...
+    def dropna(self, how: Literal["any", "all"] = ...): ...
     def repeat(self, repeats, axis=...): ...
     def take(
         self, indices, axis: int = ..., allow_fill: bool = ..., fill_value=..., **kwargs

pandas-stubs/core/indexes/multi.pyi

@@ -1,6 +1,7 @@
 from typing import (
     Callable,
     Hashable,
+    Literal,
     Sequence,
 )
 
@@ -83,7 +84,7 @@ class MultiIndex(Index):
     def is_monotonic_decreasing(self) -> bool: ...
     def duplicated(self, keep: str = ...): ...
     def fillna(self, value=..., downcast=...) -> None: ...
-    def dropna(self, how: str = ...): ...
+    def dropna(self, how: Literal["any", "all"] = ...) -> MultiIndex: ...
     def get_value(self, series, key): ...
     def get_level_values(self, level: str | int) -> Index: ...
     def unique(self, level=...): ...

pandas-stubs/core/series.pyi

@@ -869,7 +869,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     def dropna(
         self,
         axis: SeriesAxisType = ...,
-        how: _str | None = ...,
+        how: Literal["any", "all"] | None = ...,
         *,
         inplace: Literal[True],
     ) -> None: ...
@@ -878,7 +878,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         self,
         axis: SeriesAxisType = ...,
         inplace: _bool = ...,
-        how: _str | None = ...,
+        how: Literal["any", "all"] | None = ...,
     ) -> Series[S1]: ...
     def to_timestamp(
         self,

tests/test_indexes.py

@@ -75,3 +75,15 @@ def test_str_split() -> None:
     ind = pd.Index(["a-b", "c-d"])
     check(assert_type(ind.str.split("-"), pd.Index), pd.Index)
     check(assert_type(ind.str.split("-", expand=True), pd.MultiIndex), pd.MultiIndex)
+
+
+def test_index_dropna():
+    idx = pd.Index([1, 2])
+
+    check(assert_type(idx.dropna(how="all"), pd.Index), pd.Index)
+    check(assert_type(idx.dropna(how="any"), pd.Index), pd.Index)
+
+    midx = pd.MultiIndex.from_arrays([[1, 2], [3, 4]])
+
+    check(assert_type(midx.dropna(how="all"), pd.MultiIndex), pd.MultiIndex)
+    check(assert_type(midx.dropna(how="any"), pd.MultiIndex), pd.MultiIndex)