TYP: Simple type fixes for ExtensionArray (#41255)

Author: Dr-Irv

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"

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]
@@ -678,7 +683,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.
 
@@ -1205,7 +1215,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.
 
@@ -1218,7 +1228,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.
 
@@ -1292,7 +1302,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__)}")
 
     # ------------------------------------------------------------------------

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,

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,