TYP: Typing changes for ExtensionArray.astype

pandas/_testing/asserters.py

@@ -399,9 +399,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:

pandas/core/arrays/base.py

@@ -17,6 +17,7 @@
     Sequence,
     TypeVar,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -26,6 +27,7 @@
     ArrayLike,
     Dtype,
     FillnaOptions,
+    NpDtype,
     PositionalIndexer,
     Shape,
 )
@@ -516,7 +518,15 @@ def nbytes(self) -> int:
     # Additional Methods
     # ------------------------------------------------------------------------
 
-    def astype(self, dtype, copy=True):
+    @overload
+    def astype(self, dtype: NpDtype, 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 with 'dtype'.
 
@@ -531,8 +541,9 @@ def astype(self, dtype, copy=True):
 
         Returns
         -------
-        array : ndarray
-            NumPy ndarray with 'dtype' for its dtype.
+        array : ArrayLike
+            An ExtensionArray if dtype StringDtype or same as that of underlying array.
+            Otherwise a NumPy ndarray with 'dtype' for its dtype.
         """
         from pandas.core.arrays.string_ import StringDtype
 
@@ -548,7 +559,11 @@ def astype(self, dtype, copy=True):
             # 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:
         """
@@ -933,7 +948,9 @@ 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
+        # error: Incompatible return value type (got "Tuple[Union[ExtensionArray,
+        # ndarray], float]", expected "Tuple[ndarray, Any]")
+        return self.astype(object), np.nan  # type: ignore[return-value]
 
     def factorize(self, na_sentinel: int = -1) -> tuple[np.ndarray, ExtensionArray]:
         """

pandas/core/arrays/boolean.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import numbers
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    overload,
+)
 import warnings
 
 import numpy as np
@@ -13,6 +16,7 @@
 from pandas._typing import (
     ArrayLike,
     Dtype,
+    NpDtype,
     type_t,
 )
 from pandas.compat.numpy import function as nv
@@ -392,7 +396,16 @@ 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: NpDtype, 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'.
 

pandas/core/arrays/categorical.py

@@ -11,6 +11,7 @@
     TypeVar,
     Union,
     cast,
+    overload,
 )
 from warnings import (
     catch_warnings,
@@ -487,6 +488,14 @@ def _constructor(self) -> type[Categorical]:
     def _from_sequence(cls, scalars, *, dtype: Dtype | None = None, copy=False):
         return Categorical(scalars, dtype=dtype, copy=copy)
 
+    @overload
+    def astype(self, dtype: NpDtype, 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

pandas/core/arrays/floating.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from typing import overload
 import warnings
 
 import numpy as np
@@ -10,7 +11,9 @@
 )
 from pandas._typing import (
     ArrayLike,
+    Dtype,
     DtypeObj,
+    NpDtype,
 )
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import cache_readonly
@@ -271,7 +274,15 @@ 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:
+    @overload
+    def astype(self, dtype: NpDtype, 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'.
 

pandas/core/arrays/integer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from typing import overload
 import warnings
 
 import numpy as np
@@ -13,6 +14,7 @@
     ArrayLike,
     Dtype,
     DtypeObj,
+    NpDtype,
 )
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import cache_readonly
@@ -335,7 +337,15 @@ 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:
+    @overload
+    def astype(self, dtype: NpDtype, 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'.
 

pandas/core/arrays/masked.py

@@ -5,6 +5,7 @@
     Any,
     Sequence,
     TypeVar,
+    overload,
 )
 
 import numpy as np
@@ -301,6 +302,14 @@ def to_numpy(  # type: ignore[override]
             data = self._data.astype(dtype, copy=copy)
         return data
 
+    @overload
+    def astype(self, dtype: NpDtype, copy: bool = True) -> np.ndarray:
+        ...
+
+    @overload
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
+        ...
+
     def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
         dtype = pandas_dtype(dtype)
 

pandas/core/groupby/groupby.py

@@ -1497,7 +1497,7 @@ def _bool_agg(self, val_test, skipna):
         Shared func to call any / all Cython GroupBy implementations.
         """
 
-        def objs_to_bool(vals: ArrayLike) -> tuple[np.ndarray, type]:
+        def objs_to_bool(vals: ArrayLike) -> tuple[ArrayLike, type]:
             if is_object_dtype(vals):
                 vals = np.array([bool(x) for x in vals])
             elif isinstance(vals, BaseMaskedArray):

pandas/core/internals/blocks.py

@@ -1349,7 +1349,9 @@ def get_values(self, dtype: DtypeObj | None = None) -> np.ndarray:
         """
         values = self.values
         if dtype == _dtype_obj:
-            values = values.astype(object)
+            # error: Incompatible types in assignment (expression has type
+            # "Union[ExtensionArray, ndarray]", variable has type "ExtensionArray")
+            values = values.astype(object)  # type: ignore[assignment]
         # TODO(EA2D): reshape not needed with 2D EAs
         return np.asarray(values).reshape(self.shape)