TYP: change ArrayLike/AnyArrayLike alias to Union (#40379)

Author: simonjayhawkins

pandas/_typing.py

@@ -47,7 +47,7 @@
     from pandas.core.dtypes.dtypes import ExtensionDtype
 
     from pandas import Interval
-    from pandas.core.arrays.base import ExtensionArray  # noqa: F401
+    from pandas.core.arrays.base import ExtensionArray
     from pandas.core.frame import DataFrame
     from pandas.core.generic import NDFrame  # noqa: F401
     from pandas.core.groupby.generic import (
@@ -74,8 +74,8 @@
 
 # array-like
 
-AnyArrayLike = TypeVar("AnyArrayLike", "ExtensionArray", "Index", "Series", np.ndarray)
-ArrayLike = TypeVar("ArrayLike", "ExtensionArray", np.ndarray)
+ArrayLike = Union["ExtensionArray", np.ndarray]
+AnyArrayLike = Union[ArrayLike, "Index", "Series"]
 
 # scalars
 

pandas/core/algorithms.py

@@ -176,9 +176,7 @@ def _ensure_data(values: ArrayLike) -> Tuple[np.ndarray, DtypeObj]:
         elif is_timedelta64_dtype(values.dtype):
             from pandas import TimedeltaIndex
 
-            # error: Incompatible types in assignment (expression has type
-            # "TimedeltaArray", variable has type "ndarray")
-            values = TimedeltaIndex(values)._data  # type: ignore[assignment]
+            values = TimedeltaIndex(values)._data
         else:
             # Datetime
             if values.ndim > 1 and is_datetime64_ns_dtype(values.dtype):
@@ -194,22 +192,13 @@ def _ensure_data(values: ArrayLike) -> Tuple[np.ndarray, DtypeObj]:
 
             from pandas import DatetimeIndex
 
-            # Incompatible types in assignment (expression has type "DatetimeArray",
-            # variable has type "ndarray")
-            values = DatetimeIndex(values)._data  # type: ignore[assignment]
+            values = DatetimeIndex(values)._data
         dtype = values.dtype
-        # error: Item "ndarray" of "Union[PeriodArray, Any, ndarray]" has no attribute
-        # "asi8"
-        return values.asi8, dtype  # type: ignore[union-attr]
+        return values.asi8, dtype
 
     elif is_categorical_dtype(values.dtype):
-        # error: Incompatible types in assignment (expression has type "Categorical",
-        # variable has type "ndarray")
-        values = cast("Categorical", values)  # type: ignore[assignment]
-        # error: Incompatible types in assignment (expression has type "ndarray",
-        # variable has type "ExtensionArray")
-        # error: Item "ndarray" of "Union[Any, ndarray]" has no attribute "codes"
-        values = values.codes  # type: ignore[assignment,union-attr]
+        values = cast("Categorical", values)
+        values = values.codes
         dtype = pandas_dtype("category")
 
         # we are actually coercing to int64
@@ -222,10 +211,7 @@ def _ensure_data(values: ArrayLike) -> Tuple[np.ndarray, DtypeObj]:
         return values, dtype  # type: ignore[return-value]
 
     # we have failed, return object
-
-    # error: Incompatible types in assignment (expression has type "ndarray", variable
-    # has type "ExtensionArray")
-    values = np.asarray(values, dtype=object)  # type: ignore[assignment]
+    values = np.asarray(values, dtype=object)
     return ensure_object(values), np.dtype("object")
 
 
@@ -335,9 +321,7 @@ def _get_values_for_rank(values: ArrayLike):
     if is_categorical_dtype(values):
         values = cast("Categorical", values)._values_for_rank()
 
-    # error: Incompatible types in assignment (expression has type "ndarray", variable
-    # has type "ExtensionArray")
-    values, _ = _ensure_data(values)  # type: ignore[assignment]
+    values, _ = _ensure_data(values)
     return values
 
 
@@ -503,42 +487,15 @@ def isin(comps: AnyArrayLike, values: AnyArrayLike) -> np.ndarray:
         )
 
     if not isinstance(values, (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray)):
-        # error: Incompatible types in assignment (expression has type "ExtensionArray",
-        # variable has type "Index")
-        # error: Incompatible types in assignment (expression has type "ExtensionArray",
-        # variable has type "Series")
-        # error: Incompatible types in assignment (expression has type "ExtensionArray",
-        # variable has type "ndarray")
-        values = _ensure_arraylike(list(values))  # type: ignore[assignment]
+        values = _ensure_arraylike(list(values))
     elif isinstance(values, ABCMultiIndex):
         # Avoid raising in extract_array
-
-        # error: Incompatible types in assignment (expression has type "ndarray",
-        # variable has type "ExtensionArray")
-        # error: Incompatible types in assignment (expression has type "ndarray",
-        # variable has type "Index")
-        # error: Incompatible types in assignment (expression has type "ndarray",
-        # variable has type "Series")
-        values = np.array(values)  # type: ignore[assignment]
+        values = np.array(values)
     else:
-        # error: Incompatible types in assignment (expression has type "Union[Any,
-        # ExtensionArray]", variable has type "Index")
-        # error: Incompatible types in assignment (expression has type "Union[Any,
-        # ExtensionArray]", variable has type "Series")
-        values = extract_array(values, extract_numpy=True)  # type: ignore[assignment]
-
-    # error: Incompatible types in assignment (expression has type "ExtensionArray",
-    # variable has type "Index")
-    # error: Incompatible types in assignment (expression has type "ExtensionArray",
-    # variable has type "Series")
-    # error: Incompatible types in assignment (expression has type "ExtensionArray",
-    # variable has type "ndarray")
-    comps = _ensure_arraylike(comps)  # type: ignore[assignment]
-    # error: Incompatible types in assignment (expression has type "Union[Any,
-    # ExtensionArray]", variable has type "Index")
-    # error: Incompatible types in assignment (expression has type "Union[Any,
-    # ExtensionArray]", variable has type "Series")
-    comps = extract_array(comps, extract_numpy=True)  # type: ignore[assignment]
+        values = extract_array(values, extract_numpy=True)
+
+    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"
@@ -1000,9 +957,7 @@ def duplicated(values: ArrayLike, keep: Union[str, bool] = "first") -> np.ndarra
     -------
     duplicated : ndarray
     """
-    # error: Incompatible types in assignment (expression has type "ndarray", variable
-    # has type "ExtensionArray")
-    values, _ = _ensure_data(values)  # type: ignore[assignment]
+    values, _ = _ensure_data(values)
     ndtype = values.dtype.name
     f = getattr(htable, f"duplicated_{ndtype}")
     return f(values, keep=keep)

pandas/core/array_algos/putmask.py

@@ -191,16 +191,10 @@ 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)
 
-        # error: Incompatible types in assignment (expression has type "ndarray",
-        # variable has type "ExtensionArray")
-        mask = mask.to_numpy(dtype=bool, na_value=False)  # type: ignore[assignment]
-
-    # error: Incompatible types in assignment (expression has type "ndarray", variable
-    # has type "ExtensionArray")
-    mask = np.asarray(mask, dtype=bool)  # type: ignore[assignment]
-    # error: Incompatible return value type (got "ExtensionArray", expected "ndarray")
-    return mask  # type: ignore[return-value]
+    mask = np.asarray(mask, dtype=bool)
+    return mask
 
 
 def setitem_datetimelike_compat(values: np.ndarray, num_set: int, other):

pandas/core/array_algos/quantile.py

@@ -40,10 +40,9 @@ def quantile_compat(values: ArrayLike, qs, interpolation: str, axis: int) -> Arr
     if isinstance(values, np.ndarray):
         fill_value = na_value_for_dtype(values.dtype, compat=False)
         mask = isna(values)
-        result = quantile_with_mask(values, mask, fill_value, qs, interpolation, axis)
+        return quantile_with_mask(values, mask, fill_value, qs, interpolation, axis)
     else:
-        result = quantile_ea_compat(values, qs, interpolation, axis)
-    return result
+        return quantile_ea_compat(values, qs, interpolation, axis)
 
 
 def quantile_with_mask(

pandas/core/array_algos/replace.py

@@ -95,9 +95,7 @@ def _check_comparison_types(
 
     if is_numeric_v_string_like(a, b):
         # GH#29553 avoid deprecation warnings from numpy
-        # error: Incompatible return value type (got "ndarray", expected
-        # "Union[ExtensionArray, bool]")
-        return np.zeros(a.shape, dtype=bool)  # type: ignore[return-value]
+        return np.zeros(a.shape, dtype=bool)
 
     elif is_datetimelike_v_numeric(a, b):
         # GH#29553 avoid deprecation warnings from numpy

pandas/core/array_algos/take.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
 
 import functools
-from typing import Optional
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+    overload,
+)
 
 import numpy as np
 
@@ -20,6 +24,33 @@
 
 from pandas.core.construction import ensure_wrapped_if_datetimelike
 
+if TYPE_CHECKING:
+    from pandas.core.arrays.base import ExtensionArray
+
+
+@overload
+def take_nd(
+    arr: np.ndarray,
+    indexer,
+    axis: int = ...,
+    out: Optional[np.ndarray] = ...,
+    fill_value=...,
+    allow_fill: bool = ...,
+) -> np.ndarray:
+    ...
+
+
+@overload
+def take_nd(
+    arr: ExtensionArray,
+    indexer,
+    axis: int = ...,
+    out: Optional[np.ndarray] = ...,
+    fill_value=...,
+    allow_fill: bool = ...,
+) -> ArrayLike:
+    ...
+
 
 def take_nd(
     arr: ArrayLike,

pandas/core/arrays/_mixins.py

@@ -291,8 +291,7 @@ def fillna(
             value, mask, len(self)  # type: ignore[arg-type]
         )
 
-        # error: "ExtensionArray" has no attribute "any"
-        if mask.any():  # type: ignore[attr-defined]
+        if mask.any():
             if method is not None:
                 # TODO: check value is None
                 # (for now) when self.ndim == 2, we assume axis=0

pandas/core/arrays/base.py

@@ -10,6 +10,7 @@
 
 import operator
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
@@ -71,6 +72,16 @@
     nargsort,
 )
 
+if TYPE_CHECKING:
+
+    class ExtensionArraySupportsAnyAll("ExtensionArray"):
+        def any(self, *, skipna: bool = True) -> bool:
+            pass
+
+        def all(self, *, skipna: bool = True) -> bool:
+            pass
+
+
 _extension_array_shared_docs: Dict[str, str] = {}
 
 ExtensionArrayT = TypeVar("ExtensionArrayT", bound="ExtensionArray")
@@ -380,7 +391,7 @@ def __iter__(self):
         for i in range(len(self)):
             yield self[i]
 
-    def __contains__(self, item) -> bool:
+    def __contains__(self, item) -> Union[bool, np.bool_]:
         """
         Return for `item in self`.
         """
@@ -391,8 +402,7 @@ def __contains__(self, item) -> bool:
             if not self._can_hold_na:
                 return False
             elif item is self.dtype.na_value or isinstance(item, self.dtype.type):
-                # error: "ExtensionArray" has no attribute "any"
-                return self.isna().any()  # type: ignore[attr-defined]
+                return self.isna().any()
             else:
                 return False
         else:
@@ -543,7 +553,7 @@ def astype(self, dtype, copy=True):
 
         return np.array(self, dtype=dtype, copy=copy)
 
-    def isna(self) -> ArrayLike:
+    def isna(self) -> Union[np.ndarray, ExtensionArraySupportsAnyAll]:
         """
         A 1-D array indicating if each value is missing.
 
@@ -648,8 +658,7 @@ def argmin(self, skipna: bool = True) -> int:
         ExtensionArray.argmax
         """
         validate_bool_kwarg(skipna, "skipna")
-        # error: "ExtensionArray" has no attribute "any"
-        if not skipna and self.isna().any():  # type: ignore[attr-defined]
+        if not skipna and self.isna().any():
             raise NotImplementedError
         return nargminmax(self, "argmin")
 
@@ -673,8 +682,7 @@ def argmax(self, skipna: bool = True) -> int:
         ExtensionArray.argmin
         """
         validate_bool_kwarg(skipna, "skipna")
-        # error: "ExtensionArray" has no attribute "any"
-        if not skipna and self.isna().any():  # type: ignore[attr-defined]
+        if not skipna and self.isna().any():
             raise NotImplementedError
         return nargminmax(self, "argmax")
 
@@ -714,8 +722,7 @@ def fillna(self, value=None, method=None, limit=None):
             value, mask, len(self)  # type: ignore[arg-type]
         )
 
-        # error: "ExtensionArray" has no attribute "any"
-        if mask.any():  # type: ignore[attr-defined]
+        if mask.any():
             if method is not None:
                 func = missing.get_fill_func(method)
                 new_values, _ = func(self.astype(object), limit=limit, mask=mask)
@@ -1156,9 +1163,7 @@ def view(self, dtype: Optional[Dtype] = None) -> ArrayLike:
         #   giving a view with the same dtype as self.
         if dtype is not None:
             raise NotImplementedError(dtype)
-        # error: Incompatible return value type (got "Union[ExtensionArray, Any]",
-        # expected "ndarray")
-        return self[:]  # type: ignore[return-value]
+        return self[:]
 
     # ------------------------------------------------------------------------
     # Printing

pandas/core/arrays/boolean.py

@@ -406,18 +406,14 @@ def astype(self, dtype, copy: bool = True) -> ArrayLike:
         dtype = pandas_dtype(dtype)
 
         if isinstance(dtype, ExtensionDtype):
-            # error: Incompatible return value type (got "ExtensionArray", expected
-            # "ndarray")
-            return super().astype(dtype, copy)  # type: ignore[return-value]
+            return super().astype(dtype, copy)
 
         if is_bool_dtype(dtype):
             # astype_nansafe converts np.nan to True
             if self._hasna:
                 raise ValueError("cannot convert float NaN to bool")
             else:
-                # error: Incompatible return value type (got "ndarray", expected
-                # "ExtensionArray")
-                return self._data.astype(dtype, copy=copy)  # type: ignore[return-value]
+                return self._data.astype(dtype, copy=copy)
 
         # for integer, error if there are missing values
         if is_integer_dtype(dtype) and self._hasna:
@@ -429,12 +425,7 @@ def astype(self, dtype, copy: bool = True) -> ArrayLike:
         if is_float_dtype(dtype):
             na_value = np.nan
         # coerce
-
-        # error: Incompatible return value type (got "ndarray", expected
-        # "ExtensionArray")
-        return self.to_numpy(  # type: ignore[return-value]
-            dtype=dtype, na_value=na_value, copy=False
-        )
+        return self.to_numpy(dtype=dtype, na_value=na_value, copy=False)
 
     def _values_for_argsort(self) -> np.ndarray:
         """

pandas/core/arrays/categorical.py

@@ -550,8 +550,7 @@ def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
                 new_cats, libalgos.ensure_platform_int(self._codes)
             )
 
-        # error: Incompatible return value type (got "Categorical", expected "ndarray")
-        return result  # type: ignore[return-value]
+        return result
 
     @cache_readonly
     def itemsize(self) -> int:
@@ -2659,8 +2658,9 @@ def _get_codes_for_values(values, categories: Index) -> np.ndarray:
     # Only hit here when we've already coerced to object dtypee.
 
     hash_klass, vals = get_data_algo(values)
-    # error: Value of type variable "ArrayLike" of "get_data_algo" cannot be "Index"
-    _, cats = get_data_algo(categories)  # type: ignore[type-var]
+    # pandas/core/arrays/categorical.py:2661: error: Argument 1 to "get_data_algo" has
+    # incompatible type "Index"; expected "Union[ExtensionArray, ndarray]"  [arg-type]
+    _, cats = get_data_algo(categories)  # type: ignore[arg-type]
     t = hash_klass(len(cats))
     t.map_locations(cats)
     return coerce_indexer_dtype(t.lookup(vals), cats)