make extension arrays work with new typing, fixing astype and to_numpy

pandas-dev · Dr-Irv · Jan 23, 2021 · Jan 23, 2021 · Jan 31, 2021 · Jan 31, 2021
commit 548c198015be1d0040f0ae6061acf8a0b246e5dd
diff --git a/pandas/_testing/asserters.py b/pandas/_testing/asserters.py
@@ -382,9 +382,13 @@ def _get_ilevel_values(index, level):
     # skip exact index checking when `check_categorical` is False
     if check_exact and check_categorical:
         if not left.equals(right):
-            diff = (
-                np.sum((left._values != right._values).astype(int)) * 100.0 / len(left)
-            )
+            #  error: Value of type variable "_Number" of "sum" cannot be
+            # "Union[ExtensionArray, ndarray, Any]"
+            thesum = np.sum(
+                (left._values != right._values).astype(int)
+            )  # type: ignore[type-var]
+            # error: Unsupported operand types for * ("ExtensionArray" and "float")
+            diff = thesum * 100.0 / len(left)  # type: ignore[operator]
             msg = f"{obj} values are different ({np.round(diff, 5)} %)"
             raise_assert_detail(obj, msg, left, right)
     else:

diff --git a/pandas/_typing.py b/pandas/_typing.py
@@ -76,11 +76,17 @@
 
 ArrayLike = Union["ExtensionArray", np.ndarray]
 AnyArrayLike = Union[ArrayLike, "Index", "Series"]
-
+AnySequenceLike = Union[
+    "ExtensionArray",
+    "Index",
+    "Series",
+    Sequence[Any],
+    np.ndarray,
+]
 # scalars
 
 PythonScalar = Union[str, int, float, bool]
-DatetimeLikeScalar = TypeVar("DatetimeLikeScalar", "Period", "Timestamp", "Timedelta")
+DatetimeLikeScalar = Union["Period", "Timestamp", "Timedelta"]
 PandasScalar = Union["Period", "Timestamp", "Timedelta", "Interval"]
 Scalar = Union[PythonScalar, PandasScalar]
 

diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py
@@ -30,7 +30,9 @@
 )
 from pandas._typing import (
     AnyArrayLike,
+    AnySequenceLike,
     ArrayLike,
+    Dtype,
     DtypeObj,
     FrameOrSeriesUnion,
 )
@@ -216,7 +218,7 @@ def _ensure_data(values: ArrayLike) -> Tuple[np.ndarray, DtypeObj]:
 
 
 def _reconstruct_data(
-    values: ArrayLike, dtype: DtypeObj, original: AnyArrayLike
+    values: ArrayLike, dtype: Dtype, original: AnyArrayLike
 ) -> ArrayLike:
     """
     reverse of _ensure_data
@@ -244,31 +246,28 @@ def _reconstruct_data(
 
         values = cls._from_sequence(values)
     elif is_bool_dtype(dtype):
-        # error: Argument 1 to "astype" of "_ArrayOrScalarCommon" has
-        # incompatible type "Union[dtype, ExtensionDtype]"; expected
-        # "Union[dtype, None, type, _SupportsDtype, str, Tuple[Any, int],
-        # Tuple[Any, Union[int, Sequence[int]]], List[Any], _DtypeDict,
-        # Tuple[Any, Any]]"
+        # error: Argument 1 to "astype" of "_ArrayOrScalarCommon" has incompatible
+        # type "Union[ExtensionDtype, Union[str, dtype[Any]], Type[str], Type[float],
+        # Type[int], Type[complex], Type[bool], Type[object]]"; expected
+        # "Union[dtype[Any], None, type, _SupportsDType, str, Union[Tuple[Any, int],
+        # Tuple[Any, Union[int, Sequence[int]]], List[Any], _DTypeDict,
+        # Tuple[Any, Any]]]"
         values = values.astype(dtype, copy=False)  # type: ignore[arg-type]
 
         # we only support object dtypes bool Index
         if isinstance(original, ABCIndex):
             values = values.astype(object, copy=False)
     elif dtype is not None:
         if is_datetime64_dtype(dtype):
-            # error: Incompatible types in assignment (expression has type
-            # "str", variable has type "Union[dtype, ExtensionDtype]")
-            dtype = "datetime64[ns]"  # type: ignore[assignment]
+            dtype = "datetime64[ns]"
         elif is_timedelta64_dtype(dtype):
-            # error: Incompatible types in assignment (expression has type
-            # "str", variable has type "Union[dtype, ExtensionDtype]")
-            dtype = "timedelta64[ns]"  # type: ignore[assignment]
-
-        # error: Argument 1 to "astype" of "_ArrayOrScalarCommon" has
-        # incompatible type "Union[dtype, ExtensionDtype]"; expected
-        # "Union[dtype, None, type, _SupportsDtype, str, Tuple[Any, int],
-        # Tuple[Any, Union[int, Sequence[int]]], List[Any], _DtypeDict,
-        # Tuple[Any, Any]]"
+            dtype = "timedelta64[ns]"
+        # error: Argument 1 to "astype" of "_ArrayOrScalarCommon" has incompatible
+        # type "Union[ExtensionDtype, Union[str, dtype[Any]], Type[str], Type[float],
+        # Type[int], Type[complex], Type[bool], Type[object]]"; expected
+        # "Union[dtype[Any], None, type, _SupportsDType, str, Union[Tuple[Any, int],
+        # Tuple[Any, Union[int, Sequence[int]]], List[Any], _DTypeDict,
+        # Tuple[Any, Any]]]"
         values = values.astype(dtype, copy=False)  # type: ignore[arg-type]
 
     return values
@@ -461,7 +460,7 @@ def unique(values):
 unique1d = unique
 
 
-def isin(comps: AnyArrayLike, values: AnyArrayLike) -> np.ndarray:
+def isin(comps: AnySequenceLike, values: AnySequenceLike) -> np.ndarray:
     """
     Compute the isin boolean array.
 
@@ -497,9 +496,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

diff --git a/pandas/core/array_algos/putmask.py b/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

diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py
@@ -3,6 +3,7 @@
 from functools import wraps
 from typing import (
     Any,
+    List,
     Optional,
     Sequence,
     Type,
@@ -260,12 +261,13 @@ def __getitem__(self: NDArrayBackedExtensionArrayT, key: int) -> Any:
 
     @overload
     def __getitem__(
-        self: NDArrayBackedExtensionArrayT, key: Union[slice, np.ndarray]
+        self: NDArrayBackedExtensionArrayT, key: Union[slice, np.ndarray, List[Any]]
     ) -> NDArrayBackedExtensionArrayT:
         ...
 
     def __getitem__(
-        self: NDArrayBackedExtensionArrayT, key: Union[int, slice, np.ndarray]
+        self: NDArrayBackedExtensionArrayT,
+        key: Union[int, slice, np.ndarray, List[Any]],
     ) -> Union[NDArrayBackedExtensionArrayT, Any]:
         if lib.is_integer(key):
             # fast-path

diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py
@@ -15,6 +15,7 @@
     Callable,
     Dict,
     Iterator,
+    List,
     Literal,
     Optional,
     Sequence,
@@ -32,6 +33,7 @@
 from pandas._typing import (
     ArrayLike,
     Dtype,
+    NpDtype,
     Shape,
 )
 from pandas.compat import set_function_name
@@ -301,11 +303,11 @@ def __getitem__(self, item: int) -> Any:
         ...
 
     @overload
-    def __getitem__(self, item: Union[slice, np.ndarray]) -> ExtensionArray:
+    def __getitem__(self, item: Union[slice, np.ndarray, List[Any]]) -> ExtensionArray:
         ...
 
     def __getitem__(
-        self, item: Union[int, slice, np.ndarray]
+        self, item: Union[int, slice, np.ndarray, List[Any]]
     ) -> Union[ExtensionArray, Any]:
         """
         Select a subset of self.
@@ -441,9 +443,10 @@ def __ne__(self, other: Any) -> ArrayLike:  # type: ignore[override]
 
     def to_numpy(
         self,
-        dtype: Optional[Dtype] = None,
+        dtype: Optional[NpDtype] = None,
         copy: bool = False,
         na_value: Optional[Any] = lib.no_default,
+        **kwargs: Any,
     ) -> np.ndarray:
         """
         Convert to a NumPy ndarray.
@@ -470,12 +473,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:
@@ -527,8 +525,15 @@ def nbytes(self) -> int:
     # ------------------------------------------------------------------------
     # Additional Methods
     # ------------------------------------------------------------------------
+    @overload
+    def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray:
+        ...
+
+    @overload
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
+        ...
 
-    def astype(self, dtype: Dtype, copy: bool = True) -> np.ndarray:
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
         """
         Cast to a NumPy array with 'dtype'.
 
@@ -562,7 +567,7 @@ def astype(self, dtype: Dtype, copy: bool = True) -> np.ndarray:
         ):  # allow conversion to StringArrays
             return dtype.construct_array_type()._from_sequence(self, copy=False)
 
-        return np.array(self, dtype=dtype, copy=copy)
+        return np.array(self, dtype=cast(NpDtype, dtype), copy=copy)
 
     def isna(self) -> Union[np.ndarray, ExtensionArraySupportsAnyAll]:
         """
@@ -829,9 +834,9 @@ def unique(self) -> ExtensionArray:
 
     def searchsorted(
         self,
-        value: ArrayLike,
+        value: Sequence[Any],
         side: Literal["left", "right"] = "left",
-        sorter: Optional[ArrayLike] = None,
+        sorter: Optional[Sequence[Any]] = None,
     ) -> np.ndarray:
         """
         Find indices where elements should be inserted to maintain order.
@@ -877,7 +882,7 @@ def searchsorted(
         # 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)
+        arr = cast(np.ndarray, self.astype(object))
         return arr.searchsorted(value, side=side, sorter=sorter)
 
     def equals(self, other: object) -> bool:
@@ -914,7 +919,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: Union[ExtensionArray, Sequence[Any]]) -> np.ndarray:
+    def isin(self, values: Sequence[Any]) -> np.ndarray:
         """
         Pointwise comparison for set containment in the given values.
 
@@ -928,7 +933,7 @@ def isin(self, values: Union[ExtensionArray, Sequence[Any]]) -> np.ndarray:
         -------
         np.ndarray[bool]
         """
-        return isin(np.asarray(self), values)
+        return isin(self.astype(object), values)
 
     def _values_for_factorize(self) -> Tuple[np.ndarray, Any]:
         """
@@ -952,7 +957,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]:
         """

diff --git a/pandas/core/arrays/boolean.py b/pandas/core/arrays/boolean.py
@@ -8,6 +8,7 @@
     Tuple,
     Type,
     Union,
+    overload,
 )
 import warnings
 
@@ -20,6 +21,7 @@
 from pandas._typing import (
     ArrayLike,
     Dtype,
+    DtypeArg,
 )
 from pandas.compat.numpy import function as nv
 
@@ -296,7 +298,7 @@ def dtype(self) -> BooleanDtype:
 
     @classmethod
     def _from_sequence(
-        cls, scalars, *, dtype: Optional[Dtype] = None, copy: bool = False
+        cls, scalars, *, dtype: Optional[DtypeArg] = None, copy: bool = False
     ) -> BooleanArray:
         if dtype:
             assert dtype == "boolean"
@@ -379,7 +381,15 @@ 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:
+    @overload
+    def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray:
+        ...
+
+    @overload
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
+        ...
+
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
         """
         Cast to a NumPy array or ExtensionArray with 'dtype'.
 

diff --git a/pandas/core/arrays/categorical.py b/pandas/core/arrays/categorical.py
@@ -15,6 +15,7 @@
     TypeVar,
     Union,
     cast,
+    overload,
 )
 from warnings import warn
 
@@ -479,6 +480,14 @@ def _constructor(self) -> Type[Categorical]:
     def _from_sequence(cls, scalars, *, dtype: Optional[Dtype] = None, copy=False):
         return Categorical(scalars, dtype=dtype, copy=copy)
 
+    @overload
+    def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray:
+        ...
+
+    @overload
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
+        ...
+
     def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
         """
         Coerce this type to another dtype
@@ -2454,11 +2463,7 @@ def _str_get_dummies(self, sep="|"):
         # sep may not be in categories. Just bail on this.
         from pandas.core.arrays import PandasArray
 
-        # error: Argument 1 to "PandasArray" has incompatible type
-        # "ExtensionArray"; expected "Union[ndarray, PandasArray]"
-        return PandasArray(self.astype(str))._str_get_dummies(  # type: ignore[arg-type]
-            sep
-        )
+        return PandasArray(self.astype(str))._str_get_dummies(sep)
 
 
 # The Series.cat accessor