diff --git a/pandas/_typing.py b/pandas/_typing.py index a58dc0dba1bf1..1e1fffdd60676 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -37,6 +37,7 @@ # https://mypy.readthedocs.io/en/latest/common_issues.html#import-cycles if TYPE_CHECKING: from typing import ( + Literal, TypedDict, final, ) @@ -189,6 +190,12 @@ str, int, Sequence[Union[str, int]], Mapping[Hashable, Union[str, int]] ] +# Arguments for fillna() +if TYPE_CHECKING: + FillnaOptions = Literal["backfill", "bfill", "ffill", "pad"] +else: + FillnaOptions = str + # internals Manager = Union[ "ArrayManager", "SingleArrayManager", "BlockManager", "SingleBlockManager" diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index bd01191719143..16d42e4a51927 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -13,6 +13,7 @@ TYPE_CHECKING, Any, Callable, + Iterator, Sequence, TypeVar, cast, @@ -24,6 +25,7 @@ from pandas._typing import ( ArrayLike, Dtype, + FillnaOptions, PositionalIndexer, Shape, ) @@ -69,6 +71,7 @@ ) if TYPE_CHECKING: + from typing import Literal class ExtensionArraySupportsAnyAll("ExtensionArray"): def any(self, *, skipna: bool = True) -> bool: @@ -375,7 +378,7 @@ def __len__(self) -> int: """ raise AbstractMethodError(self) - def __iter__(self): + def __iter__(self) -> Iterator[Any]: """ Iterate over elements of the array. """ @@ -385,7 +388,7 @@ def __iter__(self): for i in range(len(self)): yield self[i] - def __contains__(self, item) -> bool | np.bool_: + def __contains__(self, item: object) -> bool | np.bool_: """ Return for `item in self`. """ @@ -400,7 +403,9 @@ def __contains__(self, item) -> bool | np.bool_: else: return False else: - return (item == self).any() + # error: Item "ExtensionArray" of "Union[ExtensionArray, ndarray]" has no + # attribute "any" + return (item == self).any() # type: ignore[union-attr] # error: Signature of "__eq__" incompatible with supertype "object" def __eq__(self, other: Any) -> ArrayLike: # type: ignore[override] @@ -680,7 +685,12 @@ def argmax(self, skipna: bool = True) -> int: raise NotImplementedError return nargminmax(self, "argmax") - def fillna(self, value=None, method=None, limit=None): + def fillna( + self, + value: object | ArrayLike | None = None, + method: FillnaOptions | None = None, + limit: int | None = None, + ): """ Fill NA/NaN values using the specified method. @@ -1207,7 +1217,7 @@ def _formatter(self, boxed: bool = False) -> Callable[[Any], str | None]: # Reshaping # ------------------------------------------------------------------------ - def transpose(self, *axes) -> ExtensionArray: + def transpose(self, *axes: int) -> ExtensionArray: """ Return a transposed view on this array. @@ -1220,7 +1230,7 @@ def transpose(self, *axes) -> ExtensionArray: def T(self) -> ExtensionArray: return self.transpose() - def ravel(self, order="C") -> ExtensionArray: + def ravel(self, order: Literal["C", "F", "A", "K"] | None = "C") -> ExtensionArray: """ Return a flattened view on this array. @@ -1294,7 +1304,7 @@ def _reduce(self, name: str, *, skipna: bool = True, **kwargs): """ raise TypeError(f"cannot perform {name} with type {self.dtype}") - def __hash__(self): + def __hash__(self) -> int: raise TypeError(f"unhashable type: {repr(type(self).__name__)}") # ------------------------------------------------------------------------ diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 50837e1b3ed50..a508b3db2f038 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -54,6 +54,7 @@ CompressionOptions, Dtype, FilePathOrBuffer, + FillnaOptions, FloatFormatType, FormattersType, FrameOrSeriesUnion, @@ -5015,7 +5016,7 @@ def rename( def fillna( self, value=..., - method: str | None = ..., + method: FillnaOptions | None = ..., axis: Axis | None = ..., inplace: Literal[False] = ..., limit=..., @@ -5027,7 +5028,7 @@ def fillna( def fillna( self, value, - method: str | None, + method: FillnaOptions | None, axis: Axis | None, inplace: Literal[True], limit=..., @@ -5060,7 +5061,7 @@ def fillna( def fillna( self, *, - method: str | None, + method: FillnaOptions | None, inplace: Literal[True], limit=..., downcast=..., @@ -5082,7 +5083,7 @@ def fillna( def fillna( self, *, - method: str | None, + method: FillnaOptions | None, axis: Axis | None, inplace: Literal[True], limit=..., @@ -5106,7 +5107,7 @@ def fillna( def fillna( self, value, - method: str | None, + method: FillnaOptions | None, *, inplace: Literal[True], limit=..., @@ -5118,7 +5119,7 @@ def fillna( def fillna( self, value=..., - method: str | None = ..., + method: FillnaOptions | None = ..., axis: Axis | None = ..., inplace: bool = ..., limit=..., @@ -5129,8 +5130,8 @@ def fillna( @doc(NDFrame.fillna, **_shared_doc_kwargs) def fillna( self, - value=None, - method: str | None = None, + value: object | ArrayLike | None = None, + method: FillnaOptions | None = None, axis: Axis | None = None, inplace: bool = False, limit=None, diff --git a/pandas/core/series.py b/pandas/core/series.py index 240f678960969..c8e9898f9462a 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -38,6 +38,7 @@ Axis, Dtype, DtypeObj, + FillnaOptions, FrameOrSeriesUnion, IndexKeyFunc, NpDtype, @@ -4594,7 +4595,7 @@ def drop( def fillna( self, value=..., - method: str | None = ..., + method: FillnaOptions | None = ..., axis: Axis | None = ..., inplace: Literal[False] = ..., limit=..., @@ -4606,7 +4607,7 @@ def fillna( def fillna( self, value, - method: str | None, + method: FillnaOptions | None, axis: Axis | None, inplace: Literal[True], limit=..., @@ -4639,7 +4640,7 @@ def fillna( def fillna( self, *, - method: str | None, + method: FillnaOptions | None, inplace: Literal[True], limit=..., downcast=..., @@ -4661,7 +4662,7 @@ def fillna( def fillna( self, *, - method: str | None, + method: FillnaOptions | None, axis: Axis | None, inplace: Literal[True], limit=..., @@ -4685,7 +4686,7 @@ def fillna( def fillna( self, value, - method: str | None, + method: FillnaOptions | None, *, inplace: Literal[True], limit=..., @@ -4697,7 +4698,7 @@ def fillna( def fillna( self, value=..., - method: str | None = ..., + method: FillnaOptions | None = ..., axis: Axis | None = ..., inplace: bool = ..., limit=..., @@ -4709,8 +4710,8 @@ def fillna( @doc(NDFrame.fillna, **_shared_doc_kwargs) # type: ignore[has-type] def fillna( self, - value=None, - method=None, + value: object | ArrayLike | None = None, + method: FillnaOptions | None = None, axis=None, inplace=False, limit=None,