From e8bc2730dea81fdd9e3732665339795a101fd234 Mon Sep 17 00:00:00 2001 From: Brock Date: Sun, 1 Nov 2020 09:15:40 -0800 Subject: [PATCH] REF: move _wrap_joined_index up to NDarrayBackedExtensionIndex --- pandas/core/indexes/category.py | 8 -------- pandas/core/indexes/datetimelike.py | 18 ++++++++++-------- pandas/core/indexes/extension.py | 9 ++++++++- 3 files changed, 18 insertions(+), 17 deletions(-) diff --git a/pandas/core/indexes/category.py b/pandas/core/indexes/category.py index c137509a2cd2d..2f2836519d847 100644 --- a/pandas/core/indexes/category.py +++ b/pandas/core/indexes/category.py @@ -29,7 +29,6 @@ from pandas.core.indexes.base import Index, _index_shared_docs, maybe_extract_name from pandas.core.indexes.extension import NDArrayBackedExtensionIndex, inherit_names import pandas.core.missing as missing -from pandas.core.ops import get_op_result_name _index_doc_kwargs = dict(ibase._index_doc_kwargs) _index_doc_kwargs.update(dict(target_klass="CategoricalIndex")) @@ -669,10 +668,3 @@ def _delegate_method(self, name: str, *args, **kwargs): if is_scalar(res): return res return CategoricalIndex(res, name=self.name) - - def _wrap_joined_index( - self, joined: np.ndarray, other: "CategoricalIndex" - ) -> "CategoricalIndex": - name = get_op_result_name(self, other) - cat = self._data._from_backing_data(joined) - return type(self)._simple_new(cat, name=name) diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index 751eafaa0d78e..9215fc8994d87 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -38,7 +38,6 @@ make_wrapped_arith_op, ) from pandas.core.indexes.numeric import Int64Index -from pandas.core.ops import get_op_result_name from pandas.core.tools.timedeltas import to_timedelta if TYPE_CHECKING: @@ -629,20 +628,23 @@ def insert(self, loc: int, item): def _can_union_without_object_cast(self, other) -> bool: return is_dtype_equal(self.dtype, other.dtype) - def _wrap_joined_index(self, joined: np.ndarray, other): - assert other.dtype == self.dtype, (other.dtype, self.dtype) - name = get_op_result_name(self, other) - + def _get_join_freq(self, other): + """ + Get the freq to attach to the result of a join operation. + """ if is_period_dtype(self.dtype): freq = self.freq else: self = cast(DatetimeTimedeltaMixin, self) freq = self.freq if self._can_fast_union(other) else None + return freq - new_data = self._data._from_backing_data(joined) - new_data._freq = freq + def _wrap_joined_index(self, joined: np.ndarray, other): + assert other.dtype == self.dtype, (other.dtype, self.dtype) - return type(self)._simple_new(new_data, name=name) + result = super()._wrap_joined_index(joined, other) + result._data._freq = self._get_join_freq(other) + return result @doc(Index._convert_arr_indexer) def _convert_arr_indexer(self, keyarr): diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index 1f26ceaf2d1b7..3376aef8bef87 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -1,7 +1,7 @@ """ Shared methods for Index subclasses backed by ExtensionArray. """ -from typing import List +from typing import List, TypeVar import numpy as np @@ -18,6 +18,8 @@ from pandas.core.indexes.base import Index from pandas.core.ops import get_op_result_name +_T = TypeVar("_T", bound="NDArrayBackedExtensionIndex") + def inherit_from_data(name: str, delegate, cache: bool = False, wrap: bool = False): """ @@ -338,3 +340,8 @@ def putmask(self, mask, value): np.putmask(new_values, mask, value) new_arr = self._data._from_backing_data(new_values) return type(self)._simple_new(new_arr, name=self.name) + + def _wrap_joined_index(self: _T, joined: np.ndarray, other: _T) -> _T: + name = get_op_result_name(self, other) + arr = self._data._from_backing_data(joined) + return type(self)._simple_new(arr, name=name)