Skip to content

Commit

Permalink
REF: tighten what we accept in TimedeltaIndex._simple_new (#31315)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored Feb 9, 2020
1 parent 79ca148 commit c988567
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 40 deletions.
5 changes: 4 additions & 1 deletion pandas/core/arrays/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,9 +195,12 @@ def __init__(self, values, dtype=_TD_DTYPE, freq=None, copy=False):
def _simple_new(cls, values, freq=None, dtype=_TD_DTYPE):
assert dtype == _TD_DTYPE, dtype
assert isinstance(values, np.ndarray), type(values)
if values.dtype != _TD_DTYPE:
assert values.dtype == "i8"
values = values.view(_TD_DTYPE)

result = object.__new__(cls)
result._data = values.view(_TD_DTYPE)
result._data = values
result._freq = to_offset(freq)
result._dtype = _TD_DTYPE
return result
Expand Down
19 changes: 16 additions & 3 deletions pandas/core/indexes/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -582,7 +582,8 @@ def delete(self, loc):
if loc.start in (0, None) or loc.stop in (len(self), None):
freq = self.freq

return self._shallow_copy(new_i8s, freq=freq)
arr = type(self._data)._simple_new(new_i8s, dtype=self.dtype, freq=freq)
return type(self)._simple_new(arr, name=self.name)


class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin, Int64Index):
Expand Down Expand Up @@ -623,6 +624,14 @@ def _shallow_copy(self, values=None, **kwargs):
if values is None:
values = self._data

if isinstance(values, type(self)):
values = values._data
if isinstance(values, np.ndarray):
# TODO: We would rather not get here
if kwargs.get("freq") is not None:
raise ValueError(kwargs)
values = type(self._data)(values, dtype=self.dtype)

attributes = self._get_attributes_dict()

if "freq" not in kwargs and self.freq is not None:
Expand Down Expand Up @@ -801,7 +810,10 @@ def _union(self, other, sort):
this, other = self._maybe_utc_convert(other)

if this._can_fast_union(other):
return this._fast_union(other, sort=sort)
result = this._fast_union(other, sort=sort)
if result.freq is None:
result._set_freq("infer")
return result
else:
i8self = Int64Index._simple_new(self.asi8, name=self.name)
i8other = Int64Index._simple_new(other.asi8, name=other.name)
Expand Down Expand Up @@ -934,7 +946,8 @@ def insert(self, loc, item):
new_i8s = np.concatenate(
(self[:loc].asi8, [item.view(np.int64)], self[loc:].asi8)
)
return self._shallow_copy(new_i8s, freq=freq)
arr = type(self._data)._simple_new(new_i8s, dtype=self.dtype, freq=freq)
return type(self)._simple_new(arr, name=self.name)
except (AttributeError, TypeError):

# fall back to object index
Expand Down
15 changes: 4 additions & 11 deletions pandas/core/indexes/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,22 +173,15 @@ def __new__(
def _simple_new(cls, values, name=None, freq=None, dtype=_TD_DTYPE):
# `dtype` is passed by _shallow_copy in corner cases, should always
# be timedelta64[ns] if present

if not isinstance(values, TimedeltaArray):
values = TimedeltaArray._simple_new(values, dtype=dtype, freq=freq)
else:
if freq is None:
freq = values.freq
assert isinstance(values, TimedeltaArray), type(values)
assert dtype == _TD_DTYPE, dtype
assert values.dtype == "m8[ns]", values.dtype
assert isinstance(values, TimedeltaArray)
assert freq is None or values.freq == freq

tdarr = TimedeltaArray._simple_new(values._data, freq=freq)
result = object.__new__(cls)
result._data = tdarr
result._data = values
result._name = name
# For groupby perf. See note in indexes/base about _index_data
result._index_data = tdarr._data
result._index_data = values._data

result._reset_identity()
return result
Expand Down
33 changes: 28 additions & 5 deletions pandas/core/resample.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from pandas.core.groupby.groupby import GroupBy, _GroupBy, _pipe_template, get_groupby
from pandas.core.groupby.grouper import Grouper
from pandas.core.groupby.ops import BinGrouper
from pandas.core.indexes.api import Index
from pandas.core.indexes.datetimes import DatetimeIndex, date_range
from pandas.core.indexes.period import PeriodIndex, period_range
from pandas.core.indexes.timedeltas import TimedeltaIndex, timedelta_range
Expand Down Expand Up @@ -424,10 +425,7 @@ def _wrap_result(self, result):

if isinstance(result, ABCSeries) and result.empty:
obj = self.obj
if isinstance(obj.index, PeriodIndex):
result.index = obj.index.asfreq(self.freq)
else:
result.index = obj.index._shallow_copy(freq=self.freq)
result.index = _asfreq_compat(obj.index, freq=self.freq)
result.name = getattr(obj, "name", None)

return result
Expand Down Expand Up @@ -1787,8 +1785,8 @@ def asfreq(obj, freq, method=None, how=None, normalize=False, fill_value=None):

elif len(obj.index) == 0:
new_obj = obj.copy()
new_obj.index = obj.index._shallow_copy(freq=to_offset(freq))

new_obj.index = _asfreq_compat(obj.index, freq)
else:
dti = date_range(obj.index[0], obj.index[-1], freq=freq)
dti.name = obj.index.name
Expand All @@ -1797,3 +1795,28 @@ def asfreq(obj, freq, method=None, how=None, normalize=False, fill_value=None):
new_obj.index = new_obj.index.normalize()

return new_obj


def _asfreq_compat(index, freq):
"""
Helper to mimic asfreq on (empty) DatetimeIndex and TimedeltaIndex.
Parameters
----------
index : PeriodIndex, DatetimeIndex, or TimedeltaIndex
freq : DateOffset
Returns
-------
same type as index
"""
if len(index) != 0:
# This should never be reached, always checked by the caller
raise ValueError(
"Can only set arbitrary freq for empty DatetimeIndex or TimedeltaIndex"
)
if isinstance(index, PeriodIndex):
new_index = index.asfreq(freq=freq)
else:
new_index = Index([], dtype=index.dtype, freq=freq, name=index.name)
return new_index
31 changes: 11 additions & 20 deletions pandas/tests/resample/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from pandas.core.indexes.datetimes import date_range
from pandas.core.indexes.period import PeriodIndex, period_range
from pandas.core.indexes.timedeltas import TimedeltaIndex, timedelta_range
from pandas.core.resample import _asfreq_compat

# a fixture value can be overridden by the test parameter value. Note that the
# value of the fixture can be overridden this way even if the test doesn't use
Expand Down Expand Up @@ -103,10 +104,8 @@ def test_resample_empty_series(freq, empty_series, resample_method):
result = getattr(s.resample(freq), resample_method)()

expected = s.copy()
if isinstance(s.index, PeriodIndex):
expected.index = s.index.asfreq(freq=freq)
else:
expected.index = s.index._shallow_copy(freq=freq)
expected.index = _asfreq_compat(s.index, freq)

tm.assert_index_equal(result.index, expected.index)
assert result.index.freq == expected.index.freq
tm.assert_series_equal(result, expected, check_dtype=False)
Expand All @@ -119,10 +118,8 @@ def test_resample_count_empty_series(freq, empty_series, resample_method):
# GH28427
result = getattr(empty_series.resample(freq), resample_method)()

if isinstance(empty_series.index, PeriodIndex):
index = empty_series.index.asfreq(freq=freq)
else:
index = empty_series.index._shallow_copy(freq=freq)
index = _asfreq_compat(empty_series.index, freq)

expected = pd.Series([], dtype="int64", index=index, name=empty_series.name)

tm.assert_series_equal(result, expected)
Expand All @@ -141,10 +138,8 @@ def test_resample_empty_dataframe(empty_frame, freq, resample_method):
# GH14962
expected = Series([], dtype=object)

if isinstance(df.index, PeriodIndex):
expected.index = df.index.asfreq(freq=freq)
else:
expected.index = df.index._shallow_copy(freq=freq)
expected.index = _asfreq_compat(df.index, freq)

tm.assert_index_equal(result.index, expected.index)
assert result.index.freq == expected.index.freq
tm.assert_almost_equal(result, expected, check_dtype=False)
Expand All @@ -162,10 +157,8 @@ def test_resample_count_empty_dataframe(freq, empty_frame):

result = empty_frame.resample(freq).count()

if isinstance(empty_frame.index, PeriodIndex):
index = empty_frame.index.asfreq(freq=freq)
else:
index = empty_frame.index._shallow_copy(freq=freq)
index = _asfreq_compat(empty_frame.index, freq)

expected = pd.DataFrame({"a": []}, dtype="int64", index=index)

tm.assert_frame_equal(result, expected)
Expand All @@ -181,10 +174,8 @@ def test_resample_size_empty_dataframe(freq, empty_frame):

result = empty_frame.resample(freq).size()

if isinstance(empty_frame.index, PeriodIndex):
index = empty_frame.index.asfreq(freq=freq)
else:
index = empty_frame.index._shallow_copy(freq=freq)
index = _asfreq_compat(empty_frame.index, freq)

expected = pd.Series([], dtype="int64", index=index)

tm.assert_series_equal(result, expected)
Expand Down

0 comments on commit c988567

Please sign in to comment.