TYP: Typing changes for ExtensionArray.astype

pandas/_testing/asserters.py

@@ -399,9 +399,8 @@ 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)
-            )
+            thesum = np.sum((left._values != right._values).astype(int))
+            diff = thesum * 100.0 / len(left)
             msg = f"{obj} values are different ({np.round(diff, 5)} %)"
             raise_assert_detail(obj, msg, left, right)
     else:

pandas/_typing.py

@@ -124,10 +124,8 @@
 Axes = Collection[Any]
 
 # dtypes
-NpDtype = Union[str, np.dtype]
-Dtype = Union[
-    "ExtensionDtype", NpDtype, type_t[Union[str, float, int, complex, bool, object]]
-]
+NpDtype = Union[str, np.dtype, type_t[Union[str, float, int, complex, bool, object]]]
+Dtype = Union["ExtensionDtype", NpDtype]
 # DtypeArg specifies all allowable dtypes in a functions its dtype argument
 DtypeArg = Union[Dtype, Dict[Hashable, Dtype]]
 DtypeObj = Union[np.dtype, "ExtensionDtype"]

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,
 )
@@ -515,9 +517,17 @@ def nbytes(self) -> int:
     # Additional Methods
     # ------------------------------------------------------------------------
 
-    def astype(self, dtype, copy=True):
+    @overload
+    def astype(self, dtype: NpDtype, copy: bool = ...) -> np.ndarray:
+        ...
+
+    @overload
+    def astype(self, dtype: Dtype, copy: bool = ...) -> ArrayLike:
+        ...
+
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
         """
-        Cast to a NumPy array with 'dtype'.
+        Cast to a NumPy array or ExtensionArray with 'dtype'.
 
         Parameters
         ----------
@@ -530,8 +540,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
 
@@ -547,7 +558,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:
         """

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
@@ -2468,11 +2477,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

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
@@ -280,9 +281,7 @@ def to_numpy(  # type: ignore[override]
         if na_value is lib.no_default:
             na_value = libmissing.NA
         if dtype is None:
-            # error: Incompatible types in assignment (expression has type
-            # "Type[object]", variable has type "Union[str, dtype[Any], None]")
-            dtype = object  # type: ignore[assignment]
+            dtype = object
         if self._hasna:
             if (
                 not is_object_dtype(dtype)
@@ -301,6 +300,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/arrays/sparse/array.py

@@ -523,9 +523,7 @@ def __array__(self, dtype: NpDtype | None = None) -> np.ndarray:
             try:
                 dtype = np.result_type(self.sp_values.dtype, type(fill_value))
             except TypeError:
-                # error: Incompatible types in assignment (expression has type
-                # "Type[object]", variable has type "Union[str, dtype[Any], None]")
-                dtype = object  # type: ignore[assignment]
+                dtype = object
 
         out = np.full(self.shape, fill_value, dtype=dtype)
         out[self.sp_index.to_int_index().indices] = self.sp_values

pandas/core/common.py

@@ -229,12 +229,7 @@ def asarray_tuplesafe(values, dtype: NpDtype | None = None) -> np.ndarray:
         # expected "ndarray")
         return values._values  # type: ignore[return-value]
 
-    # error: Non-overlapping container check (element type: "Union[str, dtype[Any],
-    # None]", container item type: "type")
-    if isinstance(values, list) and dtype in [  # type: ignore[comparison-overlap]
-        np.object_,
-        object,
-    ]:
+    if isinstance(values, list) and dtype in [np.object_, object]:
         return construct_1d_object_array_from_listlike(values)
 
     result = np.asarray(values, dtype=dtype)

pandas/core/construction.py

@@ -588,9 +588,7 @@ def _sanitize_ndim(
         if is_object_dtype(dtype) and isinstance(dtype, ExtensionDtype):
             # i.e. PandasDtype("O")
 
-            # error: Argument "dtype" to "asarray_tuplesafe" has incompatible type
-            # "Type[object]"; expected "Union[str, dtype[Any], None]"
-            result = com.asarray_tuplesafe(data, dtype=object)  # type: ignore[arg-type]
+            result = com.asarray_tuplesafe(data, dtype=object)
             cls = dtype.construct_array_type()
             result = cls._from_sequence(result, dtype=dtype)
         else:

pandas/core/dtypes/dtypes.py

@@ -29,7 +29,6 @@
 from pandas._typing import (
     Dtype,
     DtypeObj,
-    NpDtype,
     Ordered,
     type_t,
 )
@@ -1291,7 +1290,7 @@ class PandasDtype(ExtensionDtype):
 
     _metadata = ("_dtype",)
 
-    def __init__(self, dtype: NpDtype | PandasDtype | None):
+    def __init__(self, dtype: str | np.dtype | PandasDtype | None):
         if isinstance(dtype, PandasDtype):
             # make constructor univalent
             dtype = dtype.numpy_dtype

pandas/core/internals/blocks.py

@@ -1348,7 +1348,7 @@ def get_values(self, dtype: DtypeObj | None = None) -> np.ndarray:
         """
         return object dtype as boxed values, such as Timestamps/Timedelta
         """
-        values = self.values
+        values: ArrayLike = self.values
         if dtype == _dtype_obj:
             values = values.astype(object)
         # TODO(EA2D): reshape not needed with 2D EAs