diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index c3c91cea43f6b..283d13c0204d7 100644 --- a/pandas/core/arrays/base.py +++ b/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. diff --git a/pandas/core/base.py b/pandas/core/base.py index 9fe1af776dd2b..c59193d826d7e 100644 --- a/pandas/core/base.py +++ b/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) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 22ba317e78e63..16e64e9ee3e3d 100644 --- a/pandas/core/indexes/base.py +++ b/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. diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index 6d2253c5dc87d..b803b02d3da74 100644 --- a/pandas/core/ops/__init__.py +++ b/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. diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index c393b8028113b..7208281b333f0 100644 --- a/pandas/core/ops/array_ops.py +++ b/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 `^`. diff --git a/pandas/core/ops/dispatch.py b/pandas/core/ops/dispatch.py index 61a3032c7a02c..5c34cb20be266 100644 --- a/pandas/core/ops/dispatch.py +++ b/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. diff --git a/pandas/io/pytables.py b/pandas/io/pytables.py index c1e12887b0150..5229bfca7a84f 100644 --- a/pandas/io/pytables.py +++ b/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. """