diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 2ae988e21..1514c720c 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/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] = ... diff --git a/pandas-stubs/core/indexes/extension.pyi b/pandas-stubs/core/indexes/extension.pyi index ee31106d2..47d32773b 100644 --- a/pandas-stubs/core/indexes/extension.pyi +++ b/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 diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index b7b797917..a96d61d61 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/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=...): ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index a91e7f8e7..cc5c5759f 100644 --- a/pandas-stubs/core/series.pyi +++ b/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, diff --git a/tests/test_indexes.py b/tests/test_indexes.py index eb3797171..ed9acd539 100644 --- a/tests/test_indexes.py +++ b/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)