Skip to content

handle NaT add/sub in one place #19903

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 27, 2018
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 40 additions & 6 deletions pandas/core/indexes/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -638,6 +638,28 @@ def _add_datelike(self, other):
def _sub_datelike(self, other):
raise com.AbstractMethodError(self)

def _add_nat(self):
"""Add pd.NaT to self"""
if is_period_dtype(self):
raise TypeError('Cannot add {cls} and {typ}'
.format(cls=type(self).__name__,
typ=type(NaT).__name__))

# GH#19124 pd.NaT is treated like a timedelta for both timedelta
# and datetime dtypes
return self._nat_new(box=True)

def _sub_nat(self):
"""Subtract pd.NaT from self"""
# GH#19124 Timedelta - datetime is not in general well-defined.
# We make an exception for pd.NaT, which in this case quacks
# like a timedelta.
# For datetime64 dtypes by convention we treat NaT as a datetime, so
# this subtraction returns a timedelta64 dtype.
# For period dtype, timedelta64 is a close-enough return dtype.
result = self._nat_new(box=False)
return result.view('timedelta64[ns]')

def _sub_period(self, other):
return NotImplemented

Expand Down Expand Up @@ -684,6 +706,8 @@ def __add__(self, other):
return NotImplemented

# scalar others
elif other is NaT:
result = self._add_nat()
elif isinstance(other, (DateOffset, timedelta, np.timedelta64)):
result = self._add_delta(other)
elif isinstance(other, (datetime, np.datetime64)):
Expand Down Expand Up @@ -718,9 +742,13 @@ def __add__(self, other):
else: # pragma: no cover
return NotImplemented

if result is not NotImplemented:
res_name = ops.get_op_result_name(self, other)
result.name = res_name
if result is NotImplemented:
return NotImplemented
elif not isinstance(result, Index):
# Index.__new__ will choose appropriate subclass for dtype
result = Index(result)
res_name = ops.get_op_result_name(self, other)
result.name = res_name
return result

cls.__add__ = __add__
Expand All @@ -738,6 +766,8 @@ def __sub__(self, other):
return NotImplemented

# scalar others
elif other is NaT:
result = self._sub_nat()
elif isinstance(other, (DateOffset, timedelta, np.timedelta64)):
result = self._add_delta(-other)
elif isinstance(other, (datetime, np.datetime64)):
Expand Down Expand Up @@ -776,9 +806,13 @@ def __sub__(self, other):
else: # pragma: no cover
return NotImplemented

if result is not NotImplemented:
res_name = ops.get_op_result_name(self, other)
result.name = res_name
if result is NotImplemented:
return NotImplemented
elif not isinstance(result, Index):
# Index.__new__ will choose appropriate subclass for dtype
result = Index(result)
res_name = ops.get_op_result_name(self, other)
result.name = res_name
return result

cls.__sub__ = __sub__
Expand Down
27 changes: 10 additions & 17 deletions pandas/core/indexes/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -853,27 +853,15 @@ def __setstate__(self, state):
raise Exception("invalid pickle state")
_unpickle_compat = __setstate__

def _add_datelike(self, other):
# adding a timedeltaindex to a datetimelike
if other is libts.NaT:
return self._nat_new(box=True)
raise TypeError("cannot add {0} and {1}"
.format(type(self).__name__,
type(other).__name__))

def _sub_datelike(self, other):
# subtract a datetime from myself, yielding a TimedeltaIndex
from pandas import TimedeltaIndex
# subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
if isinstance(other, DatetimeIndex):
# require tz compat
if not self._has_same_tz(other):
raise TypeError("DatetimeIndex subtraction must have the same "
"timezones or no timezones")
result = self._sub_datelike_dti(other)
elif isinstance(other, (datetime, np.datetime64)):
assert other is not libts.NaT
other = Timestamp(other)
if other is libts.NaT:
result = self._nat_new(box=False)
return self - libts.NaT
# require tz compat
elif not self._has_same_tz(other):
raise TypeError("Timestamp subtraction must have the same "
Expand All @@ -887,20 +875,25 @@ def _sub_datelike(self, other):
else:
raise TypeError("cannot subtract DatetimeIndex and {typ}"
.format(typ=type(other).__name__))
return TimedeltaIndex(result)
return result.view('timedelta64[ns]')

def _sub_datelike_dti(self, other):
"""subtraction of two DatetimeIndexes"""
if not len(self) == len(other):
raise ValueError("cannot add indices of unequal length")
elif not self._has_same_tz(other):
# require tz compat
raise TypeError("{cls} subtraction must have the same "
"timezones or no timezones"
.format(cls=type(self).__name__))

self_i8 = self.asi8
other_i8 = other.asi8
new_values = self_i8 - other_i8
if self.hasnans or other.hasnans:
mask = (self._isnan) | (other._isnan)
new_values[mask] = libts.iNaT
return new_values.view('i8')
return new_values.view('timedelta64[ns]')

def _maybe_update_attributes(self, attrs):
""" Update Index attributes (e.g. freq) depending on op """
Expand Down
17 changes: 1 addition & 16 deletions pandas/core/indexes/period.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
is_scalar,
is_datetime64_dtype,
is_datetime64_any_dtype,
is_timedelta64_dtype,
is_period_dtype,
is_bool_dtype,
pandas_dtype,
Expand All @@ -23,7 +22,6 @@
import pandas.tseries.frequencies as frequencies
from pandas.tseries.frequencies import get_freq_code as _gfc
from pandas.core.indexes.datetimes import DatetimeIndex, Int64Index, Index
from pandas.core.indexes.timedeltas import TimedeltaIndex
from pandas.core.indexes.datetimelike import DatelikeOps, DatetimeIndexOpsMixin
from pandas.core.tools.datetimes import parse_time_string
import pandas.tseries.offsets as offsets
Expand Down Expand Up @@ -700,16 +698,6 @@ def _maybe_convert_timedelta(self, other):
return other.n
msg = _DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr)
raise IncompatibleFrequency(msg)
elif isinstance(other, np.ndarray):
if is_integer_dtype(other):
return other
elif is_timedelta64_dtype(other):
offset = frequencies.to_offset(self.freq)
if isinstance(offset, offsets.Tick):
nanos = delta_to_nanoseconds(other)
offset_nanos = delta_to_nanoseconds(offset)
if (nanos % offset_nanos).all() == 0:
return nanos // offset_nanos
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This block is never reached b/c there is a separate check for np.ndarray before this.

elif is_integer(other):
# integer is passed to .shift via
# _add_datetimelike_methods basically
Expand All @@ -724,10 +712,7 @@ def _add_delta(self, other):
return self.shift(ordinal_delta)

def _sub_datelike(self, other):
if other is tslib.NaT:
new_data = np.empty(len(self), dtype=np.int64)
new_data.fill(tslib.iNaT)
return TimedeltaIndex(new_data)
assert other is not tslib.NaT
return NotImplemented

def _sub_period(self, other):
Expand Down
26 changes: 9 additions & 17 deletions pandas/core/indexes/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -416,25 +416,17 @@ def _evaluate_with_timedelta_like(self, other, op):
def _add_datelike(self, other):
# adding a timedeltaindex to a datetimelike
from pandas import Timestamp, DatetimeIndex
if other is NaT:
# GH#19124 pd.NaT is treated like a timedelta
return self._nat_new()
else:
other = Timestamp(other)
i8 = self.asi8
result = checked_add_with_arr(i8, other.value,
arr_mask=self._isnan)
result = self._maybe_mask_results(result, fill_value=iNaT)
return DatetimeIndex(result)
assert other is not NaT
other = Timestamp(other)
i8 = self.asi8
result = checked_add_with_arr(i8, other.value,
arr_mask=self._isnan)
result = self._maybe_mask_results(result, fill_value=iNaT)
return DatetimeIndex(result)

def _sub_datelike(self, other):
# GH#19124 Timedelta - datetime is not in general well-defined.
# We make an exception for pd.NaT, which in this case quacks
# like a timedelta.
if other is NaT:
return self._nat_new()
else:
raise TypeError("cannot subtract a datelike from a TimedeltaIndex")
assert other is not NaT
raise TypeError("cannot subtract a datelike from a TimedeltaIndex")

def _addsub_offset_array(self, other, op):
# Add or subtract Array-like of DateOffset objects
Expand Down