Fix typing for extension arrays and extension dtypes without isin and astype

pandas/core/algorithms.py

@@ -485,9 +485,11 @@ def isin(comps: AnyArrayLike, values: AnyArrayLike) -> np.ndarray:
     comps = _ensure_arraylike(comps)
     comps = extract_array(comps, extract_numpy=True)
     if is_extension_array_dtype(comps.dtype):
-        # error: Incompatible return value type (got "Series", expected "ndarray")
-        # error: Item "ndarray" of "Union[Any, ndarray]" has no attribute "isin"
-        return comps.isin(values)  # type: ignore[return-value,union-attr]
+        #  error: Argument 1 to "isin" of "ExtensionArray" has incompatible type
+        # "Union[Any, ExtensionArray, ndarray]"; expected "Sequence[Any]"
+        # error: Item "ndarray" of "Union[Any, ExtensionArray, ndarray]" has no
+        # attribute "isin"
+        return comps.isin(values)  # type: ignore[arg-type, union-attr]
 
     elif needs_i8_conversion(comps.dtype):
         # Dispatch to DatetimeLikeArrayMixin.isin

pandas/core/array_algos/putmask.py

@@ -191,7 +191,7 @@ def extract_bool_array(mask: ArrayLike) -> np.ndarray:
         # We could have BooleanArray, Sparse[bool], ...
         #  Except for BooleanArray, this is equivalent to just
         #  np.asarray(mask, dtype=bool)
-        mask = mask.to_numpy(dtype=bool, na_value=False)
+        mask = mask.to_numpy(dtype=np.dtype(bool), na_value=False)
 
     mask = np.asarray(mask, dtype=bool)
     return mask

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,
+    NpDtype,
     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.
         """
@@ -424,9 +427,9 @@ def __ne__(self, other: Any) -> ArrayLike:  # type: ignore[override]
 
     def to_numpy(
         self,
-        dtype: Dtype | None = None,
+        dtype: NpDtype | None = None,
         copy: bool = False,
-        na_value=lib.no_default,
+        na_value: Any | None = lib.no_default,
     ) -> np.ndarray:
         """
         Convert to a NumPy ndarray.
@@ -453,12 +456,7 @@ def to_numpy(
         -------
         numpy.ndarray
         """
-        # error: Argument "dtype" to "asarray" has incompatible type
-        # "Union[ExtensionDtype, str, dtype[Any], Type[str], Type[float], Type[int],
-        # Type[complex], Type[bool], Type[object], None]"; expected "Union[dtype[Any],
-        # None, type, _SupportsDType, str, Union[Tuple[Any, int], Tuple[Any, Union[int,
-        # Sequence[int]]], List[Any], _DTypeDict, Tuple[Any, Any]]]"
-        result = np.asarray(self, dtype=dtype)  # type: ignore[arg-type]
+        result = np.asarray(self, dtype=dtype)
         if copy or na_value is not lib.no_default:
             result = result.copy()
         if na_value is not lib.no_default:
@@ -510,8 +508,7 @@ def nbytes(self) -> int:
     # ------------------------------------------------------------------------
     # Additional Methods
     # ------------------------------------------------------------------------
-
-    def astype(self, dtype, copy=True):
+    def astype(self, dtype: Dtype, copy: bool = True):
         """
         Cast to a NumPy array with 'dtype'.
 
@@ -544,8 +541,11 @@ def astype(self, dtype, copy=True):
             dtype, (ArrowStringDtype, StringDtype)
         ):  # allow conversion to StringArrays
             return dtype.construct_array_type()._from_sequence(self, copy=False)
-
-        return np.array(self, dtype=dtype, copy=copy)
+        # error: Argument "dtype" to "array" has incompatible type
+        # "Union[ExtensionDtype, dtype[Any]]"; expected "Union[dtype[Any], None,
+        # type, _SupportsDType, str, Union[Tuple[Any, int], Tuple[Any,
+        # Union[int, Sequence[int]]], List[Any], _DTypeDict, Tuple[Any, Any]]]"
+        return np.array(self, dtype=dtype, copy=copy)  # type: ignore[arg-type]
 
     def isna(self) -> np.ndarray | ExtensionArraySupportsAnyAll:
         """
@@ -591,8 +591,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.
@@ -680,7 +680,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: Any | ArrayLike | None = None,
+        method: Literal["backfill", "bfill", "ffill", "pad"] | None = None,
+        limit: int | None = None,
+    ) -> ExtensionArray:
         """
         Fill NA/NaN values using the specified method.
 
@@ -729,16 +734,19 @@ 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.
 
         Returns
         -------
         valid : ExtensionArray
         """
-        # error: Unsupported operand type for ~ ("ExtensionArray")
-        return self[~self.isna()]  # type: ignore[operator]
+        # error: Incompatible return value type (got "Union[ExtensionArray, Any]",
+        # expected "ExtensionArrayT")
+        # error: Unsupported operand type for ~ ("Union[ndarray,
+        # ExtensionArraySupportsAnyAll]")
+        return self[~self.isna()]  # type: ignore[return-value, operator]
 
     def shift(self, periods: int = 1, fill_value: object = None) -> ExtensionArray:
         """
@@ -794,7 +802,7 @@ def shift(self, periods: int = 1, fill_value: object = None) -> ExtensionArray:
             b = empty
         return self._concat_same_type([a, b])
 
-    def unique(self):
+    def unique(self) -> ExtensionArray:
         """
         Compute the ExtensionArray of unique values.
 
@@ -805,7 +813,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: Sequence[Any],
+        side: Literal["left", "right"] = "left",
+        sorter: Sequence[Any] | None = None,
+    ) -> np.ndarray:
         """
         Find indices where elements should be inserted to maintain order.
 
@@ -850,7 +863,8 @@ def searchsorted(self, value, side="left", sorter=None):
         # 1. Values outside the range of the `data_for_sorting` fixture
         # 2. Values between the values in the `data_for_sorting` fixture
         # 3. Missing values.
-        arr = self.astype(object)
+        # TODO: overload astype so that cast is unnecessary
+        arr = cast(np.ndarray, self.astype(object))
         return arr.searchsorted(value, side=side, sorter=sorter)
 
     def equals(self, other: object) -> bool:
@@ -887,7 +901,7 @@ def equals(self, other: object) -> bool:
             equal_na = self.isna() & other.isna()  # type: ignore[operator]
             return bool((equal_values | equal_na).all())
 
-    def isin(self, values) -> np.ndarray:
+    def isin(self, values: Sequence[Any]) -> np.ndarray:
         """
         Pointwise comparison for set containment in the given values.
 
@@ -901,7 +915,9 @@ def isin(self, values) -> np.ndarray:
         -------
         np.ndarray[bool]
         """
-        return isin(np.asarray(self), values)
+        # error: Argument 2 to "isin" has incompatible type "Sequence[Any]"; expected
+        # "Union[Union[ExtensionArray, ndarray], Index, Series]"
+        return isin(self.astype(object), values)  # type: ignore[arg-type]
 
     def _values_for_factorize(self) -> tuple[np.ndarray, Any]:
         """
@@ -925,7 +941,7 @@ def _values_for_factorize(self) -> tuple[np.ndarray, Any]:
         The values returned by this method are also used in
         :func:`pandas.util.hash_pandas_object`.
         """
-        return self.astype(object), np.nan
+        return cast(np.ndarray, self.astype(object)), np.nan
 
     def factorize(self, na_sentinel: int = -1) -> tuple[np.ndarray, ExtensionArray]:
         """
@@ -1023,7 +1039,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,
+        repeats: int | Sequence[int],
+        axis: Literal[None, 0] = None,
+    ) -> ExtensionArray:
         nv.validate_repeat((), {"axis": axis})
         ind = np.arange(len(self)).repeat(repeats)
         return self.take(ind)
@@ -1033,12 +1053,12 @@ def repeat(self, repeats, axis=None):
     # ------------------------------------------------------------------------
 
     def take(
-        self: ExtensionArrayT,
+        self,
         indices: Sequence[int],
         *,
         allow_fill: bool = False,
         fill_value: Any = None,
-    ) -> ExtensionArrayT:
+    ) -> ExtensionArray:
         """
         Take elements from an array.
 
@@ -1127,7 +1147,7 @@ def take(self, indices, allow_fill=False, fill_value=None):
         # pandas.api.extensions.take
         raise AbstractMethodError(self)
 
-    def copy(self: ExtensionArrayT) -> ExtensionArrayT:
+    def copy(self) -> ExtensionArray:
         """
         Return a copy of the array.
 
@@ -1207,7 +1227,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 +1240,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,13 +1314,13 @@ 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, loc: int | Sequence[int]) -> ExtensionArray:
         indexer = np.delete(np.arange(len(self)), loc)
         return self.take(indexer)
 

pandas/core/arrays/boolean.py

@@ -392,7 +392,7 @@ def reconstruct(x):
     def _coerce_to_array(self, value) -> tuple[np.ndarray, np.ndarray]:
         return coerce_to_array(value)
 
-    def astype(self, dtype, copy: bool = True) -> ArrayLike:
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
         """
         Cast to a NumPy array or ExtensionArray with 'dtype'.
 

pandas/core/arrays/floating.py

@@ -10,6 +10,7 @@
 )
 from pandas._typing import (
     ArrayLike,
+    Dtype,
     DtypeObj,
 )
 from pandas.compat.numpy import function as nv
@@ -271,7 +272,7 @@ def _from_sequence_of_strings(
     def _coerce_to_array(self, value) -> tuple[np.ndarray, np.ndarray]:
         return coerce_to_array(value, dtype=self.dtype)
 
-    def astype(self, dtype, copy: bool = True) -> ArrayLike:
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
         """
         Cast to a NumPy array or ExtensionArray with 'dtype'.
 

pandas/core/arrays/interval.py

@@ -1517,7 +1517,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: 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/arrays/masked.py

@@ -208,10 +208,7 @@ def __len__(self) -> int:
     def __invert__(self: BaseMaskedArrayT) -> BaseMaskedArrayT:
         return type(self)(~self._data, self._mask.copy())
 
-    # error: Argument 1 of "to_numpy" is incompatible with supertype "ExtensionArray";
-    # supertype defines the argument type as "Union[ExtensionDtype, str, dtype[Any],
-    # Type[str], Type[float], Type[int], Type[complex], Type[bool], Type[object], None]"
-    def to_numpy(  # type: ignore[override]
+    def to_numpy(
         self,
         dtype: NpDtype | None = None,
         copy: bool = False,

pandas/core/arrays/numpy_.py

@@ -366,10 +366,7 @@ def skew(
     # ------------------------------------------------------------------------
     # Additional Methods
 
-    # error: Argument 1 of "to_numpy" is incompatible with supertype "ExtensionArray";
-    # supertype defines the argument type as "Union[ExtensionDtype, str, dtype[Any],
-    # Type[str], Type[float], Type[int], Type[complex], Type[bool], Type[object], None]"
-    def to_numpy(  # type: ignore[override]
+    def to_numpy(
         self,
         dtype: NpDtype | None = None,
         copy: bool = False,

pandas/core/arrays/string_arrow.py

@@ -15,7 +15,7 @@
     missing as libmissing,
 )
 from pandas._typing import (
-    Dtype,
+    DtypeArg,
     NpDtype,
     PositionalIndexer,
     type_t,
@@ -227,15 +227,15 @@ def _chk_pyarrow_available(cls) -> None:
             raise ImportError(msg)
 
     @classmethod
-    def _from_sequence(cls, scalars, dtype: Dtype | None = None, copy: bool = False):
+    def _from_sequence(cls, scalars, dtype: DtypeArg | None = None, copy: bool = False):
         cls._chk_pyarrow_available()
         # convert non-na-likes to str, and nan-likes to ArrowStringDtype.na_value
         scalars = lib.ensure_string_array(scalars, copy=False)
         return cls(pa.array(scalars, type=pa.string(), from_pandas=True))
 
     @classmethod
     def _from_sequence_of_strings(
-        cls, strings, dtype: Dtype | None = None, copy: bool = False
+        cls, strings, dtype: DtypeArg | None = None, copy: bool = False
     ):
         return cls._from_sequence(strings, dtype=dtype, copy=copy)
 
@@ -254,10 +254,7 @@ def __arrow_array__(self, type=None):
         """Convert myself to a pyarrow Array or ChunkedArray."""
         return self._data
 
-    # error: Argument 1 of "to_numpy" is incompatible with supertype "ExtensionArray";
-    # supertype defines the argument type as "Union[ExtensionDtype, str, dtype[Any],
-    # Type[str], Type[float], Type[int], Type[complex], Type[bool], Type[object], None]"
-    def to_numpy(  # type: ignore[override]
+    def to_numpy(
         self,
         dtype: NpDtype | None = None,
         copy: bool = False,
@@ -696,7 +693,7 @@ def value_counts(self, dropna: bool = True) -> Series:
 
     _str_na_value = ArrowStringDtype.na_value
 
-    def _str_map(self, f, na_value=None, dtype: Dtype | None = None):
+    def _str_map(self, f, na_value=None, dtype: DtypeArg | None = None):
         # TODO: de-duplicate with StringArray method. This method is moreless copy and
         # paste.