TYP: Update type hints for ExtensionArray and ExtensionDtype

pandas/core/arrays/_mixins.py

@@ -8,6 +8,7 @@
     Type,
     TypeVar,
     Union,
+    overload,
 )
 
 import numpy as np
@@ -248,6 +249,22 @@ def __setitem__(self, key, value):
     def _validate_setitem_value(self, value):
         return value
 
+    @overload
+    def __getitem__(self: NDArrayBackedExtensionArrayT, key: int) -> Any:
+        ...
+
+    @overload
+    def __getitem__(
+        self: NDArrayBackedExtensionArrayT, key: slice
+    ) -> NDArrayBackedExtensionArrayT:
+        ...
+
+    @overload
+    def __getitem__(
+        self: NDArrayBackedExtensionArrayT, key: np.ndarray
+    ) -> NDArrayBackedExtensionArrayT:
+        ...
+
     def __getitem__(
         self: NDArrayBackedExtensionArrayT, key: Union[int, slice, np.ndarray]
     ) -> Union[NDArrayBackedExtensionArrayT, Any]:

pandas/core/arrays/base.py

@@ -10,16 +10,19 @@
 
 import operator
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
+    Iterator,
     Optional,
     Sequence,
     Tuple,
     Type,
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -71,6 +74,9 @@
     nargsort,
 )
 
+if TYPE_CHECKING:
+    from typing import Literal
+
 _extension_array_shared_docs: Dict[str, str] = {}
 
 ExtensionArrayT = TypeVar("ExtensionArrayT", bound="ExtensionArray")
@@ -282,8 +288,20 @@ def _from_factorized(cls, values, original):
     # Must be a Sequence
     # ------------------------------------------------------------------------
 
+    @overload
+    def __getitem__(self: ExtensionArrayT, item: int) -> Any:
+        ...
+
+    @overload
+    def __getitem__(self: ExtensionArrayT, item: slice) -> ExtensionArrayT:
+        ...
+
+    @overload
+    def __getitem__(self: ExtensionArrayT, item: np.ndarray) -> ExtensionArrayT:
+        ...
+
     def __getitem__(
-        self, item: Union[int, slice, np.ndarray]
+        self: ExtensionArrayT, item: Union[int, slice, np.ndarray]
     ) -> Union[ExtensionArray, Any]:
         """
         Select a subset of self.
@@ -370,7 +388,7 @@ def __len__(self) -> int:
         """
         raise AbstractMethodError(self)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Any]:
         """
         Iterate over elements of the array.
         """
@@ -380,7 +398,7 @@ def __iter__(self):
         for i in range(len(self)):
             yield self[i]
 
-    def __contains__(self, item) -> bool:
+    def __contains__(self, item: Any) -> bool:
         """
         Return for `item in self`.
         """
@@ -419,7 +437,7 @@ def to_numpy(
         self,
         dtype: Optional[Dtype] = None,
         copy: bool = False,
-        na_value=lib.no_default,
+        na_value: Optional[Any] = lib.no_default,
     ) -> np.ndarray:
         """
         Convert to a NumPy ndarray.
@@ -498,7 +516,7 @@ def nbytes(self) -> int:
     # Additional Methods
     # ------------------------------------------------------------------------
 
-    def astype(self, dtype, copy=True):
+    def astype(self, dtype: Dtype, copy: bool = True) -> np.ndarray:
         """
         Cast to a NumPy array with 'dtype'.
 
@@ -578,8 +596,8 @@ def argsort(
         ascending: bool = True,
         kind: str = "quicksort",
         na_position: str = "last",
-        *args,
-        **kwargs,
+        *args: Any,
+        **kwargs: Any,
     ) -> np.ndarray:
         """
         Return the indices that would sort this array.
@@ -667,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: ExtensionArrayT,
+        value: Optional[Union[Any, ArrayLike]] = None,
+        method: Optional[Literal["backfill", "bfill", "ffill", "pad"]] = None,
+        limit: Optional[int] = None,
+    ) -> ExtensionArrayT:
         """
         Fill NA/NaN values using the specified method.
 
@@ -712,7 +735,7 @@ def fillna(self, value=None, method=None, limit=None):
             new_values = self.copy()
         return new_values
 
-    def dropna(self):
+    def dropna(self) -> ExtensionArrayT:
         """
         Return ExtensionArray without NA values.
 
@@ -722,7 +745,9 @@ def dropna(self):
         """
         return self[~self.isna()]
 
-    def shift(self, periods: int = 1, fill_value: object = None) -> ExtensionArray:
+    def shift(
+        self: ExtensionArrayT, periods: int = 1, fill_value: object = None
+    ) -> ExtensionArrayT:
         """
         Shift values by desired number.
 
@@ -770,13 +795,13 @@ def shift(self, periods: int = 1, fill_value: object = None) -> ExtensionArray:
         )
         if periods > 0:
             a = empty
-            b = self[:-periods]
+            b: ExtensionArrayT = self[:-periods]
         else:
             a = self[abs(periods) :]
             b = empty
         return self._concat_same_type([a, b])
 
-    def unique(self):
+    def unique(self: ExtensionArrayT) -> ExtensionArrayT:
         """
         Compute the ExtensionArray of unique values.
 
@@ -787,7 +812,12 @@ def unique(self):
         uniques = unique(self.astype(object))
         return self._from_sequence(uniques, dtype=self.dtype)
 
-    def searchsorted(self, value, side="left", sorter=None):
+    def searchsorted(
+        self,
+        value: ArrayLike,
+        side: Optional[Literal["left", "right"]] = "left",
+        sorter: Optional[ArrayLike] = None,
+    ) -> np.ndarray:
         """
         Find indices where elements should be inserted to maintain order.
 
@@ -868,7 +898,7 @@ def equals(self, other: object) -> bool:
             equal_na = self.isna() & other.isna()
             return bool((equal_values | equal_na).all())
 
-    def isin(self, values) -> np.ndarray:
+    def isin(self, values: Union[ExtensionArray, Sequence[Any]]) -> np.ndarray:
         """
         Pointwise comparison for set containment in the given values.
 
@@ -1002,7 +1032,11 @@ def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, ExtensionArray]:
 
     @Substitution(klass="ExtensionArray")
     @Appender(_extension_array_shared_docs["repeat"])
-    def repeat(self, repeats, axis=None):
+    def repeat(
+        self: ExtensionArrayT,
+        repeats: Union[int, Sequence[int]],
+        axis: Literal[None] = None,
+    ) -> ExtensionArrayT:
         nv.validate_repeat((), {"axis": axis})
         ind = np.arange(len(self)).repeat(repeats)
         return self.take(ind)
@@ -1186,7 +1220,7 @@ def _formatter(self, boxed: bool = False) -> Callable[[Any], Optional[str]]:
     # Reshaping
     # ------------------------------------------------------------------------
 
-    def transpose(self, *axes) -> ExtensionArray:
+    def transpose(self: ExtensionArrayT, *axes: int) -> ExtensionArrayT:
         """
         Return a transposed view on this array.
 
@@ -1196,10 +1230,12 @@ def transpose(self, *axes) -> ExtensionArray:
         return self[:]
 
     @property
-    def T(self) -> ExtensionArray:
+    def T(self: ExtensionArrayT) -> ExtensionArrayT:
         return self.transpose()
 
-    def ravel(self, order="C") -> ExtensionArray:
+    def ravel(
+        self: ExtensionArrayT, order: Optional[Literal["C", "F", "A", "K"]] = "C"
+    ) -> ExtensionArrayT:
         """
         Return a flattened view on this array.
 
@@ -1221,7 +1257,7 @@ def ravel(self, order="C") -> ExtensionArray:
     @classmethod
     def _concat_same_type(
         cls: Type[ExtensionArrayT], to_concat: Sequence[ExtensionArrayT]
-    ) -> ExtensionArrayT:
+    ):
         """
         Concatenate multiple array of this dtype.
 
@@ -1273,13 +1309,15 @@ 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__)}")
 
     # ------------------------------------------------------------------------
     # Non-Optimized Default Methods
 
-    def delete(self: ExtensionArrayT, loc) -> ExtensionArrayT:
+    def delete(
+        self: ExtensionArrayT, loc: Union[int, Sequence[int]]
+    ) -> ExtensionArrayT:
         indexer = np.delete(np.arange(len(self)), loc)
         return self.take(indexer)
 

pandas/core/arrays/interval.py

@@ -11,6 +11,7 @@
     Sequence,
     Type,
     TypeVar,
+    Union,
     cast,
 )
 
@@ -1493,7 +1494,9 @@ def delete(self: IntervalArrayT, loc) -> IntervalArrayT:
         return self._shallow_copy(left=new_left, right=new_right)
 
     @Appender(_extension_array_shared_docs["repeat"] % _shared_docs_kwargs)
-    def repeat(self: IntervalArrayT, repeats: int, axis=None) -> IntervalArrayT:
+    def repeat(
+        self: IntervalArrayT, repeats: Union[int, Sequence[int]], axis=None
+    ) -> IntervalArrayT:
         nv.validate_repeat((), {"axis": axis})
         left_repeat = self.left.repeat(repeats)
         right_repeat = self.right.repeat(repeats)

pandas/core/dtypes/base.py

@@ -12,6 +12,7 @@
     Tuple,
     Type,
     Union,
+    cast,
 )
 
 import numpy as np
@@ -211,7 +212,7 @@ def construct_array_type(cls) -> Type[ExtensionArray]:
         raise NotImplementedError
 
     @classmethod
-    def construct_from_string(cls, string: str):
+    def construct_from_string(cls, string: str) -> ExtensionDtype:
         r"""
         Construct this type from a string.
 
@@ -422,9 +423,7 @@ def register(self, dtype: Type[ExtensionDtype]) -> None:
 
         self.dtypes.append(dtype)
 
-    def find(
-        self, dtype: Union[Type[ExtensionDtype], str]
-    ) -> Optional[Type[ExtensionDtype]]:
+    def find(self, dtype: Union[Type[ExtensionDtype], str]) -> Optional[ExtensionDtype]:
         """
         Parameters
         ----------
@@ -439,7 +438,7 @@ def find(
             if not isinstance(dtype, type):
                 dtype_type = type(dtype)
             if issubclass(dtype_type, ExtensionDtype):
-                return dtype
+                return cast(ExtensionDtype, dtype)
 
             return None
 
@@ -452,4 +451,4 @@ def find(
         return None
 
 
-registry = Registry()
+registry: Registry = Registry()