diff --git a/pandas/_libs/tslibs/offsets.pyx b/pandas/_libs/tslibs/offsets.pyx index 6f5ad2ae45f50..251af50ab12ce 100644 --- a/pandas/_libs/tslibs/offsets.pyx +++ b/pandas/_libs/tslibs/offsets.pyx @@ -17,8 +17,6 @@ np.import_array() from util cimport is_string_object, is_integer_object -from pandas._libs.tslib import monthrange - from conversion cimport tz_convert_single, pydt_to_i8 from frequencies cimport get_freq_code from nattype cimport NPY_NAT @@ -471,6 +469,160 @@ cdef inline int month_add_months(pandas_datetimestruct dts, int months) nogil: return 12 if new_month == 0 else new_month +@cython.wraparound(False) +@cython.boundscheck(False) +def shift_quarters(int64_t[:] dtindex, int quarters, + int q1start_month, object day, int modby=3): + """ + Given an int64 array representing nanosecond timestamps, shift all elements + by the specified number of quarters using DateOffset semantics. + + Parameters + ---------- + dtindex : int64_t[:] timestamps for input dates + quarters : int number of quarters to shift + q1start_month : int month in which Q1 begins by convention + day : {'start', 'end', 'business_start', 'business_end'} + modby : int (3 for quarters, 12 for years) + + Returns + ------- + out : ndarray[int64_t] + """ + cdef: + Py_ssize_t i + pandas_datetimestruct dts + int count = len(dtindex) + int months_to_roll, months_since, n, compare_day + bint roll_check + int64_t[:] out = np.empty(count, dtype='int64') + + if day == 'start': + with nogil: + for i in range(count): + if dtindex[i] == NPY_NAT: + out[i] = NPY_NAT + continue + + dt64_to_dtstruct(dtindex[i], &dts) + n = quarters + + months_since = (dts.month - q1start_month) % modby + + # offset semantics - if on the anchor point and going backwards + # shift to next + if n <= 0 and (months_since != 0 or + (months_since == 0 and dts.day > 1)): + n += 1 + + dts.year = year_add_months(dts, modby * n - months_since) + dts.month = month_add_months(dts, modby * n - months_since) + dts.day = 1 + + out[i] = dtstruct_to_dt64(&dts) + + elif day == 'end': + with nogil: + for i in range(count): + if dtindex[i] == NPY_NAT: + out[i] = NPY_NAT + continue + + dt64_to_dtstruct(dtindex[i], &dts) + n = quarters + + months_since = (dts.month - q1start_month) % modby + + if n <= 0 and months_since != 0: + # The general case of this condition would be + # `months_since != 0 or (months_since == 0 and + # dts.day > get_days_in_month(dts.year, dts.month))` + # but the get_days_in_month inequality would never hold. + n += 1 + elif n > 0 and (months_since == 0 and + dts.day < get_days_in_month(dts.year, + dts.month)): + n -= 1 + + dts.year = year_add_months(dts, modby * n - months_since) + dts.month = month_add_months(dts, modby * n - months_since) + dts.day = get_days_in_month(dts.year, dts.month) + + out[i] = dtstruct_to_dt64(&dts) + + elif day == 'business_start': + with nogil: + for i in range(count): + if dtindex[i] == NPY_NAT: + out[i] = NPY_NAT + continue + + dt64_to_dtstruct(dtindex[i], &dts) + n = quarters + + months_since = (dts.month - q1start_month) % modby + compare_month = dts.month - months_since + compare_month = compare_month or 12 + # compare_day is only relevant for comparison in the case + # where months_since == 0. + compare_day = get_firstbday(dts.year, compare_month) + + if n <= 0 and (months_since != 0 or + (months_since == 0 and dts.day > compare_day)): + # make sure to roll forward, so negate + n += 1 + elif n > 0 and (months_since == 0 and dts.day < compare_day): + # pretend to roll back if on same month but + # before compare_day + n -= 1 + + dts.year = year_add_months(dts, modby * n - months_since) + dts.month = month_add_months(dts, modby * n - months_since) + + dts.day = get_firstbday(dts.year, dts.month) + + out[i] = dtstruct_to_dt64(&dts) + + elif day == 'business_end': + with nogil: + for i in range(count): + if dtindex[i] == NPY_NAT: + out[i] = NPY_NAT + continue + + dt64_to_dtstruct(dtindex[i], &dts) + n = quarters + + months_since = (dts.month - q1start_month) % modby + compare_month = dts.month - months_since + compare_month = compare_month or 12 + # compare_day is only relevant for comparison in the case + # where months_since == 0. + compare_day = get_lastbday(dts.year, compare_month) + + if n <= 0 and (months_since != 0 or + (months_since == 0 and dts.day > compare_day)): + # make sure to roll forward, so negate + n += 1 + elif n > 0 and (months_since == 0 and dts.day < compare_day): + # pretend to roll back if on same month but + # before compare_day + n -= 1 + + dts.year = year_add_months(dts, modby * n - months_since) + dts.month = month_add_months(dts, modby * n - months_since) + + dts.day = get_lastbday(dts.year, dts.month) + + out[i] = dtstruct_to_dt64(&dts) + + else: + raise ValueError("day must be None, 'start', 'end', " + "'business_start', or 'business_end'") + + return np.asarray(out) + + @cython.wraparound(False) @cython.boundscheck(False) def shift_months(int64_t[:] dtindex, int months, object day=None): diff --git a/pandas/tests/tseries/offsets/test_yqm_offsets.py b/pandas/tests/tseries/offsets/test_yqm_offsets.py index 292dd5eba938e..22b8cf6119d18 100644 --- a/pandas/tests/tseries/offsets/test_yqm_offsets.py +++ b/pandas/tests/tseries/offsets/test_yqm_offsets.py @@ -33,9 +33,15 @@ def test_quarterly_dont_normalize(): assert (result.time() == date.time()) -@pytest.mark.parametrize('offset', [MonthBegin(), MonthEnd(), - BMonthBegin(), BMonthEnd()]) -def test_apply_index(offset): +@pytest.mark.parametrize('n', [-2, 1]) +@pytest.mark.parametrize('cls', [MonthBegin, MonthEnd, + BMonthBegin, BMonthEnd, + QuarterBegin, QuarterEnd, + BQuarterBegin, BQuarterEnd, + YearBegin, YearEnd, + BYearBegin, BYearEnd]) +def test_apply_index(cls, n): + offset = cls(n=n) rng = pd.date_range(start='1/1/2000', periods=100000, freq='T') ser = pd.Series(rng) diff --git a/pandas/tseries/offsets.py b/pandas/tseries/offsets.py index 8e1ead5dfbe9e..a3cddaa19dc17 100644 --- a/pandas/tseries/offsets.py +++ b/pandas/tseries/offsets.py @@ -27,7 +27,7 @@ apply_index_wraps, roll_yearday, shift_month, - BeginMixin, EndMixin, + EndMixin, BaseOffset) @@ -1028,10 +1028,7 @@ def cbday(self): @cache_readonly def m_offset(self): - kwds = self.kwds - kwds = {key: kwds[key] for key in kwds - if key not in ['calendar', 'weekmask', 'holidays', 'offset']} - return MonthEnd(n=1, normalize=self.normalize, **kwds) + return MonthEnd(n=1, normalize=self.normalize) @apply_wraps def apply(self, other): @@ -1106,10 +1103,7 @@ def cbday(self): @cache_readonly def m_offset(self): - kwds = self.kwds - kwds = {key: kwds[key] for key in kwds - if key not in ['calendar', 'weekmask', 'holidays', 'offset']} - return MonthBegin(n=1, normalize=self.normalize, **kwds) + return MonthBegin(n=1, normalize=self.normalize) @apply_wraps def apply(self, other): @@ -1254,12 +1248,9 @@ def onOffset(self, dt): def _apply(self, n, other): # if other.day is not day_of_month move to day_of_month and update n - if other.day < self.day_of_month: - other = other.replace(day=self.day_of_month) - if n > 0: - n -= 1 + if n > 0 and other.day < self.day_of_month: + n -= 1 elif other.day > self.day_of_month: - other = other.replace(day=self.day_of_month) n += 1 months = n // 2 @@ -1309,12 +1300,9 @@ def onOffset(self, dt): def _apply(self, n, other): # if other.day is not day_of_month move to day_of_month and update n if other.day < self.day_of_month: - other = other.replace(day=self.day_of_month) n -= 1 - elif other.day > self.day_of_month: - other = other.replace(day=self.day_of_month) - if n <= 0: - n += 1 + elif n <= 0 and other.day > self.day_of_month: + n += 1 months = n // 2 + n % 2 day = 1 if n % 2 else self.day_of_month @@ -1471,6 +1459,7 @@ def apply(self, other): def getOffsetOfMonth(self, dt): w = Week(weekday=self.weekday) d = datetime(dt.year, dt.month, 1, tzinfo=dt.tzinfo) + # TODO: Is this DST-safe? d = w.rollforward(d) return d + timedelta(weeks=self.week) @@ -1550,6 +1539,7 @@ def getOffsetOfMonth(self, dt): d = datetime(dt.year, dt.month, 1, dt.hour, dt.minute, dt.second, dt.microsecond, tzinfo=dt.tzinfo) eom = m.rollforward(d) + # TODO: Is this DST-safe? w = Week(weekday=self.weekday) return w.rollback(eom) @@ -1635,6 +1625,12 @@ def onOffset(self, dt): modMonth = (dt.month - self.startingMonth) % 3 return modMonth == 0 and dt.day == self._get_offset_day(dt) + @apply_index_wraps + def apply_index(self, dtindex): + shifted = liboffsets.shift_quarters(dtindex.asi8, self.n, + self.startingMonth, self._day_opt) + return dtindex._shallow_copy(shifted) + class BQuarterEnd(QuarterOffset): """DateOffset increments between business Quarter dates @@ -1659,7 +1655,7 @@ class BQuarterBegin(QuarterOffset): _day_opt = 'business_start' -class QuarterEnd(EndMixin, QuarterOffset): +class QuarterEnd(QuarterOffset): """DateOffset increments between business Quarter dates startingMonth = 1 corresponds to dates like 1/31/2007, 4/30/2007, ... startingMonth = 2 corresponds to dates like 2/28/2007, 5/31/2007, ... @@ -1670,25 +1666,14 @@ class QuarterEnd(EndMixin, QuarterOffset): _prefix = 'Q' _day_opt = 'end' - @apply_index_wraps - def apply_index(self, i): - return self._end_apply_index(i, self.freqstr) - -class QuarterBegin(BeginMixin, QuarterOffset): +class QuarterBegin(QuarterOffset): _outputName = 'QuarterBegin' _default_startingMonth = 3 _from_name_startingMonth = 1 _prefix = 'QS' _day_opt = 'start' - @apply_index_wraps - def apply_index(self, i): - freq_month = 12 if self.startingMonth == 1 else self.startingMonth - 1 - month = liboffsets._int_to_month[freq_month] - freqstr = 'Q-{month}'.format(month=month) - return self._beg_apply_index(i, freqstr) - # --------------------------------------------------------------------- # Year-Based Offset Classes @@ -1709,6 +1694,13 @@ def apply(self, other): months = years * 12 + (self.month - other.month) return shift_month(other, months, self._day_opt) + @apply_index_wraps + def apply_index(self, dtindex): + shifted = liboffsets.shift_quarters(dtindex.asi8, self.n, + self.month, self._day_opt, + modby=12) + return dtindex._shallow_copy(shifted) + def onOffset(self, dt): if self.normalize and not _is_normalized(dt): return False @@ -1752,31 +1744,19 @@ class BYearBegin(YearOffset): _day_opt = 'business_start' -class YearEnd(EndMixin, YearOffset): +class YearEnd(YearOffset): """DateOffset increments between calendar year ends""" _default_month = 12 _prefix = 'A' _day_opt = 'end' - @apply_index_wraps - def apply_index(self, i): - # convert month anchor to annual period tuple - return self._end_apply_index(i, self.freqstr) - -class YearBegin(BeginMixin, YearOffset): +class YearBegin(YearOffset): """DateOffset increments between calendar year begin dates""" _default_month = 1 _prefix = 'AS' _day_opt = 'start' - @apply_index_wraps - def apply_index(self, i): - freq_month = 12 if self.month == 1 else self.month - 1 - month = liboffsets._int_to_month[freq_month] - freqstr = 'A-{month}'.format(month=month) - return self._beg_apply_index(i, freqstr) - # --------------------------------------------------------------------- # Special Offset Classes @@ -2245,7 +2225,8 @@ def __eq__(self, other): if isinstance(other, Tick): return self.delta == other.delta else: - return DateOffset.__eq__(self, other) + # TODO: Are there cases where this should raise TypeError? + return False # This is identical to DateOffset.__hash__, but has to be redefined here # for Python 3, because we've redefined __eq__. @@ -2261,7 +2242,8 @@ def __ne__(self, other): if isinstance(other, Tick): return self.delta != other.delta else: - return DateOffset.__ne__(self, other) + # TODO: Are there cases where this should raise TypeError? + return True @property def delta(self):