diff --git a/pandas/core/arrays/interval.py b/pandas/core/arrays/interval.py index efb66c9a47a97..757cea2c710b2 100644 --- a/pandas/core/arrays/interval.py +++ b/pandas/core/arrays/interval.py @@ -44,7 +44,11 @@ from pandas.core.arrays.base import ExtensionArray, _extension_array_shared_docs from pandas.core.arrays.categorical import Categorical import pandas.core.common as com -from pandas.core.construction import array, extract_array +from pandas.core.construction import ( + array, + ensure_wrapped_if_datetimelike, + extract_array, +) from pandas.core.indexers import check_array_indexer from pandas.core.indexes.base import ensure_index from pandas.core.ops import invalid_comparison, unpack_zerodim_and_defer @@ -251,11 +255,9 @@ def _simple_new( raise ValueError(msg) # For dt64/td64 we want DatetimeArray/TimedeltaArray instead of ndarray - from pandas.core.ops.array_ops import maybe_upcast_datetimelike_array - - left = maybe_upcast_datetimelike_array(left) + left = ensure_wrapped_if_datetimelike(left) left = extract_array(left, extract_numpy=True) - right = maybe_upcast_datetimelike_array(right) + right = ensure_wrapped_if_datetimelike(right) right = extract_array(right, extract_numpy=True) lbase = getattr(left, "_ndarray", left).base diff --git a/pandas/core/construction.py b/pandas/core/construction.py index f9ebe3f1e185e..96cf1be7520fb 100644 --- a/pandas/core/construction.py +++ b/pandas/core/construction.py @@ -402,6 +402,24 @@ def extract_array(obj: object, extract_numpy: bool = False) -> Union[Any, ArrayL return obj +def ensure_wrapped_if_datetimelike(arr): + """ + Wrap datetime64 and timedelta64 ndarrays in DatetimeArray/TimedeltaArray. + """ + if isinstance(arr, np.ndarray): + if arr.dtype.kind == "M": + from pandas.core.arrays import DatetimeArray + + return DatetimeArray._from_sequence(arr) + + elif arr.dtype.kind == "m": + from pandas.core.arrays import TimedeltaArray + + return TimedeltaArray._from_sequence(arr) + + return arr + + def sanitize_array( data, index: Optional[Index], diff --git a/pandas/core/dtypes/concat.py b/pandas/core/dtypes/concat.py index 63e3440558c75..a9355e30cd3c2 100644 --- a/pandas/core/dtypes/concat.py +++ b/pandas/core/dtypes/concat.py @@ -18,7 +18,7 @@ from pandas.core.arrays import ExtensionArray from pandas.core.arrays.sparse import SparseArray -from pandas.core.construction import array +from pandas.core.construction import array, ensure_wrapped_if_datetimelike def _get_dtype_kinds(arrays) -> Set[str]: @@ -360,12 +360,14 @@ def _concat_datetime(to_concat, axis=0): ------- a single array, preserving the combined dtypes """ - to_concat = [_wrap_datetimelike(x) for x in to_concat] + to_concat = [ensure_wrapped_if_datetimelike(x) for x in to_concat] + single_dtype = len({x.dtype for x in to_concat}) == 1 # multiple types, need to coerce to object if not single_dtype: - # wrap_datetimelike ensures that astype(object) wraps in Timestamp/Timedelta + # ensure_wrapped_if_datetimelike ensures that astype(object) wraps + # in Timestamp/Timedelta return _concatenate_2d([x.astype(object) for x in to_concat], axis=axis) if axis == 1: @@ -379,17 +381,3 @@ def _concat_datetime(to_concat, axis=0): assert result.shape[0] == 1 result = result[0] return result - - -def _wrap_datetimelike(arr): - """ - Wrap datetime64 and timedelta64 ndarrays in DatetimeArray/TimedeltaArray. - - DTA/TDA handle .astype(object) correctly. - """ - from pandas.core.construction import array as pd_array, extract_array - - arr = extract_array(arr, extract_numpy=True) - if isinstance(arr, np.ndarray) and arr.dtype.kind in ["m", "M"]: - arr = pd_array(arr) - return arr diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index c855687552e82..41d539564d91e 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -30,6 +30,7 @@ from pandas.core.dtypes.generic import ABCExtensionArray, ABCIndexClass, ABCSeries from pandas.core.dtypes.missing import isna, notna +from pandas.core.construction import ensure_wrapped_if_datetimelike from pandas.core.ops import missing from pandas.core.ops.dispatch import should_extension_dispatch from pandas.core.ops.invalid import invalid_comparison @@ -175,8 +176,8 @@ def arithmetic_op(left: ArrayLike, right: Any, op): # NB: We assume that extract_array has already been called # on `left` and `right`. - lvalues = maybe_upcast_datetimelike_array(left) - rvalues = maybe_upcast_datetimelike_array(right) + lvalues = ensure_wrapped_if_datetimelike(left) + rvalues = ensure_wrapped_if_datetimelike(right) rvalues = _maybe_upcast_for_op(rvalues, lvalues.shape) if should_extension_dispatch(lvalues, rvalues) or isinstance(rvalues, Timedelta): @@ -206,7 +207,7 @@ def comparison_op(left: ArrayLike, right: Any, op) -> ArrayLike: ndarray or ExtensionArray """ # NB: We assume extract_array has already been called on left and right - lvalues = maybe_upcast_datetimelike_array(left) + lvalues = ensure_wrapped_if_datetimelike(left) rvalues = right rvalues = lib.item_from_zerodim(rvalues) @@ -331,7 +332,7 @@ def fill_bool(x, left=None): right = construct_1d_object_array_from_listlike(right) # NB: We assume extract_array has already been called on left and right - lvalues = maybe_upcast_datetimelike_array(left) + lvalues = ensure_wrapped_if_datetimelike(left) rvalues = right if should_extension_dispatch(lvalues, rvalues): @@ -400,31 +401,6 @@ def get_array_op(op): raise NotImplementedError(op_name) -def maybe_upcast_datetimelike_array(obj: ArrayLike) -> ArrayLike: - """ - If we have an ndarray that is either datetime64 or timedelta64, wrap in EA. - - Parameters - ---------- - obj : ndarray or ExtensionArray - - Returns - ------- - ndarray or ExtensionArray - """ - if isinstance(obj, np.ndarray): - if obj.dtype.kind == "m": - from pandas.core.arrays import TimedeltaArray - - return TimedeltaArray._from_sequence(obj) - if obj.dtype.kind == "M": - from pandas.core.arrays import DatetimeArray - - return DatetimeArray._from_sequence(obj) - - return obj - - def _maybe_upcast_for_op(obj, shape: Shape): """ Cast non-pandas objects to pandas types to unify behavior of arithmetic