TYP: Arraylike

pandas/core/arrays/base.py

@@ -451,7 +451,7 @@ def isna(self) -> ArrayLike:
 
         Returns
         -------
-        na_values : Union[np.ndarray, ExtensionArray]
+        na_values : ArrayLike
             In most cases, this should return a NumPy ndarray. For
             exceptional cases like ``SparseArray``, where returning
             an ndarray would be expensive, an ExtensionArray may be
@@ -922,7 +922,7 @@ def copy(self) -> ABCExtensionArray:
         """
         raise AbstractMethodError(self)
 
-    def view(self, dtype=None) -> Union[ABCExtensionArray, np.ndarray]:
+    def view(self, dtype=None) -> ArrayLike:
         """
         Return a view on the array.
 
@@ -1175,7 +1175,7 @@ def _create_method(cls, op, coerce_to_dtype=True):
 
         Returns
         -------
-        Callable[[Any, Any], Union[ndarray, ExtensionArray]]
+        Callable[[Any, Any], ArrayLike]
             A method that can be bound to a class. When used, the method
             receives the two arguments, one of which is the instance of
             this class, and should return an ExtensionArray or an ndarray.

pandas/core/base.py

@@ -4,12 +4,12 @@
 
 import builtins
 import textwrap
-from typing import Dict, FrozenSet, List, Optional, Union
+from typing import Dict, FrozenSet, List, Optional
 
 import numpy as np
 
 import pandas._libs.lib as lib
-from pandas._typing import T
+from pandas._typing import ArrayLike, T
 from pandas.compat import PYPY
 from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
@@ -606,7 +606,7 @@ class IndexOpsMixin:
     )
 
     @property
-    def _values(self) -> Union[ExtensionArray, np.ndarray]:
+    def _values(self) -> ArrayLike:
         # must be defined here as a property for mypy
         raise AbstractMethodError(self)
 

pandas/core/indexes/base.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 import operator
 from textwrap import dedent
-from typing import Any, FrozenSet, Hashable, Optional, Union
+from typing import Any, FrozenSet, Hashable, Optional
 import warnings
 
 import numpy as np
@@ -12,7 +12,7 @@
 from pandas._libs.tslibs import OutOfBoundsDatetime, Timestamp
 from pandas._libs.tslibs.period import IncompatibleFrequency
 from pandas._libs.tslibs.timezones import tz_compare
-from pandas._typing import Label
+from pandas._typing import ArrayLike, Label
 from pandas.compat import set_function_name
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import Appender, Substitution, cache_readonly
@@ -243,7 +243,7 @@ def _outer_indexer(self, left, right):
         return libjoin.outer_join_indexer(left, right)
 
     _typ = "index"
-    _data: Union[ExtensionArray, np.ndarray]
+    _data: ArrayLike
     _id = None
     _name: Label = None
     # MultiIndex.levels previously allowed setting the index name. We
@@ -3891,7 +3891,7 @@ def array(self) -> ExtensionArray:
         return array
 
     @property
-    def _values(self) -> Union[ExtensionArray, np.ndarray]:
+    def _values(self) -> ArrayLike:
         """
         The best array representation.
 

pandas/core/ops/__init__.py

@@ -5,22 +5,17 @@
 """
 import datetime
 import operator
-from typing import Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Set, Tuple
 
 import numpy as np
 
 from pandas._libs import Timedelta, Timestamp, lib
 from pandas._libs.ops_dispatch import maybe_dispatch_ufunc_to_dunder_op  # noqa:F401
-from pandas._typing import Level
+from pandas._typing import ArrayLike, Level
 from pandas.util._decorators import Appender
 
 from pandas.core.dtypes.common import is_list_like, is_timedelta64_dtype
-from pandas.core.dtypes.generic import (
-    ABCDataFrame,
-    ABCExtensionArray,
-    ABCIndexClass,
-    ABCSeries,
-)
+from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
 from pandas.core.dtypes.missing import isna
 
 from pandas.core.construction import extract_array
@@ -61,6 +56,9 @@
     rxor,
 )
 
+if TYPE_CHECKING:
+    from pandas import Index, Series  # noqa: F401
+
 # -----------------------------------------------------------------------------
 # constants
 ARITHMETIC_BINOPS: Set[str] = {
@@ -449,12 +447,7 @@ def _align_method_SERIES(left, right, align_asobject=False):
     return left, right
 
 
-def _construct_result(
-    left: ABCSeries,
-    result: Union[np.ndarray, ABCExtensionArray],
-    index: ABCIndexClass,
-    name,
-):
+def _construct_result(left: "Series", result: ArrayLike, index: "Index", name):
     """
     Construct an appropriately-labelled Series from the result of an op.
 

pandas/core/ops/array_ops.py

@@ -4,11 +4,12 @@
 """
 from functools import partial
 import operator
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import numpy as np
 
 from pandas._libs import Timestamp, lib, ops as libops
+from pandas._typing import ArrayLike
 
 from pandas.core.dtypes.cast import (
     construct_1d_object_array_from_listlike,
@@ -155,9 +156,7 @@ def na_arithmetic_op(left, right, op, str_rep: str):
     return missing.dispatch_fill_zeros(op, left, right, result)
 
 
-def arithmetic_op(
-    left: Union[np.ndarray, ABCExtensionArray], right: Any, op, str_rep: str
-):
+def arithmetic_op(left: ArrayLike, right: Any, op, str_rep: str):
     """
     Evaluate an arithmetic operation `+`, `-`, `*`, `/`, `//`, `%`, `**`, ...
 
@@ -200,9 +199,7 @@ def arithmetic_op(
     return res_values
 
 
-def comparison_op(
-    left: Union[np.ndarray, ABCExtensionArray], right: Any, op
-) -> Union[np.ndarray, ABCExtensionArray]:
+def comparison_op(left: ArrayLike, right: Any, op) -> ArrayLike:
     """
     Evaluate a comparison operation `=`, `!=`, `>=`, `>`, `<=`, or `<`.
 
@@ -303,9 +300,7 @@ def na_logical_op(x: np.ndarray, y, op):
     return result.reshape(x.shape)
 
 
-def logical_op(
-    left: Union[np.ndarray, ABCExtensionArray], right: Any, op
-) -> Union[np.ndarray, ABCExtensionArray]:
+def logical_op(left: ArrayLike, right: Any, op) -> ArrayLike:
     """
     Evaluate a logical operation `|`, `&`, or `^`.
 

pandas/core/ops/dispatch.py

@@ -1,10 +1,12 @@
 """
 Functions for defining unary operations.
 """
-from typing import Any, Union
+from typing import Any
 
 import numpy as np
 
+from pandas._typing import ArrayLike
+
 from pandas.core.dtypes.common import (
     is_datetime64_dtype,
     is_extension_array_dtype,
@@ -13,7 +15,7 @@
     is_scalar,
     is_timedelta64_dtype,
 )
-from pandas.core.dtypes.generic import ABCExtensionArray, ABCSeries
+from pandas.core.dtypes.generic import ABCSeries
 
 from pandas.core.construction import array
 
@@ -93,9 +95,7 @@ def should_series_dispatch(left, right, op):
     return False
 
 
-def dispatch_to_extension_op(
-    op, left: Union[ABCExtensionArray, np.ndarray], right: Any,
-):
+def dispatch_to_extension_op(op, left: ArrayLike, right: Any):
     """
     Assume that left or right is a Series backed by an ExtensionArray,
     apply the operator defined by op.

pandas/io/pytables.py

@@ -2216,7 +2216,7 @@ def __eq__(self, other: Any) -> bool:
             for a in ["name", "cname", "dtype", "pos"]
         )
 
-    def set_data(self, data: Union[np.ndarray, ABCExtensionArray]):
+    def set_data(self, data: ArrayLike):
         assert data is not None
         assert self.dtype is None
 
@@ -2231,7 +2231,7 @@ def take_data(self):
         return self.data
 
     @classmethod
-    def _get_atom(cls, values: Union[np.ndarray, ABCExtensionArray]) -> "Col":
+    def _get_atom(cls, values: ArrayLike) -> "Col":
         """
         Get an appropriately typed and shaped pytables.Col object for values.
         """
@@ -4977,7 +4977,7 @@ def _dtype_to_kind(dtype_str: str) -> str:
     return kind
 
 
-def _get_data_and_dtype_name(data: Union[np.ndarray, ABCExtensionArray]):
+def _get_data_and_dtype_name(data: ArrayLike):
     """
     Convert the passed data into a storable form and a dtype string.
     """