diff --git a/pandas/_libs/tslibs/offsets.pyx b/pandas/_libs/tslibs/offsets.pyx index 933e7ed64b837..950b4d4c14772 100644 --- a/pandas/_libs/tslibs/offsets.pyx +++ b/pandas/_libs/tslibs/offsets.pyx @@ -523,11 +523,9 @@ def shift_quarters(int64_t[:] dtindex, int quarters, 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) + compare_day = get_firstbday(dts.year, dts.month) if n <= 0 and (months_since != 0 or (months_since == 0 and dts.day > compare_day)): @@ -556,11 +554,9 @@ def shift_quarters(int64_t[:] dtindex, int quarters, 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) + compare_day = get_lastbday(dts.year, dts.month) if n <= 0 and (months_since != 0 or (months_since == 0 and dts.day > compare_day)): @@ -587,15 +583,17 @@ def shift_quarters(int64_t[:] dtindex, int quarters, @cython.wraparound(False) @cython.boundscheck(False) -def shift_months(int64_t[:] dtindex, int months, object day=None): +def shift_months(int64_t[:] dtindex, int months, object day): """ Given an int64-based datetime index, shift all elements specified number of months using DateOffset semantics - day: {None, 'start', 'end'} + day: {None, 'start', 'end', 'business_start', 'business_end'} * None: day of month * 'start' 1st day of month * 'end' last day of month + * 'business_start' first business day of month + * 'business_end' last business day of month """ cdef: Py_ssize_t i @@ -721,7 +719,7 @@ def shift_months(int64_t[:] dtindex, int months, object day=None): return np.asarray(out) -cpdef datetime shift_month(datetime stamp, int months, object day_opt=None): +cpdef datetime shift_month(datetime stamp, int months, object day_opt): """ Given a datetime (or Timestamp) `stamp`, an integer `months` and an option `day_opt`, return a new datetimelike that many months later, @@ -827,7 +825,8 @@ cpdef int get_day_of_month(datetime other, day_opt) except? -1: raise ValueError(day_opt) -cpdef int roll_yearday(other, n, month, day_opt='start') except? -1: +cpdef int roll_yearday(datetime other, int n, int month, + object day_opt) except? -1: """ Possibly increment or decrement the number of periods to shift based on rollforward/rollbackward conventions. @@ -836,9 +835,12 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1: ---------- other : datetime or Timestamp n : number of periods to increment, before adjusting for rolling - day_opt : 'start', 'end' - 'start': returns 1 - 'end': returns last day of the month + month : reference month giving the first month of the year + day_opt : 'start', 'end', 'business_start', 'business_end' + 'start': compare with 1 + 'end': compare with last day of the month + 'business_start': compare with first business day of the month + 'business_end': compare with last business day of the month Returns ------- @@ -846,7 +848,7 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1: Notes ----- - * Mirrors `roll_check` in tslib.shift_months + * Mirrors `roll_check` in shift_months Examples ------- @@ -888,7 +890,7 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1: other.day < get_day_of_month(other, day_opt)): n -= 1 - elif n <= 0: + else: if other.month > month or (other.month == month and other.day > get_day_of_month(other, day_opt)): diff --git a/pandas/tests/indexes/datetimes/test_ops.py b/pandas/tests/indexes/datetimes/test_ops.py index 41d0dd38cd5f6..c2c7ab55f19f0 100644 --- a/pandas/tests/indexes/datetimes/test_ops.py +++ b/pandas/tests/indexes/datetimes/test_ops.py @@ -670,7 +670,7 @@ def test_shift_months(years, months): Timestamp('2000-01-01'), Timestamp('2000-02-29'), Timestamp('2000-12-31')]) - actual = DatetimeIndex(shift_months(s.asi8, years * 12 + months)) + actual = DatetimeIndex(shift_months(s.asi8, years * 12 + months, None)) expected = DatetimeIndex([x + pd.offsets.DateOffset( years=years, months=months) for x in s]) tm.assert_index_equal(actual, expected) diff --git a/pandas/tseries/offsets.py b/pandas/tseries/offsets.py index 54250bbf903a4..06f24cbe173c4 100644 --- a/pandas/tseries/offsets.py +++ b/pandas/tseries/offsets.py @@ -22,7 +22,7 @@ from pandas._libs.tslibs.offsets import ( ApplyTypeError, as_datetime, _is_normalized, - _get_calendar, _to_dt64, _validate_business_time, + _get_calendar, _to_dt64, _determine_offset, apply_index_wraps, roll_yearday, @@ -254,7 +254,7 @@ def apply_index(self, i): months = ((self.kwds.get('years', 0) * 12 + self.kwds.get('months', 0)) * self.n) if months: - shifted = liboffsets.shift_months(i.asi8, months) + shifted = liboffsets.shift_months(i.asi8, months, day=None) i = i._shallow_copy(shifted) weeks = (self.kwds.get('weeks', 0)) * self.n @@ -557,28 +557,31 @@ def get_str(td): def apply(self, other): if isinstance(other, datetime): n = self.n + wday = other.weekday() - if n == 0 and other.weekday() > 4: - n = 1 - - result = other - - # avoid slowness below - if abs(n) > 5: - k = n // 5 - result = result + timedelta(7 * k) - if n < 0 and result.weekday() > 4: - n += 1 - n -= 5 * k - if n == 0 and result.weekday() > 4: - n -= 1 + # avoid slowness below by operating on weeks first + weeks = n // 5 + if n <= 0 and wday > 4: + # roll forward + n += 1 - while n != 0: - k = n // abs(n) - result = result + timedelta(k) - if result.weekday() < 5: - n -= k + n -= 5 * weeks + + # n is always >= 0 at this point + if n == 0 and wday > 4: + # roll back + days = 4 - wday + elif wday > 4: + # roll forward + days = (7 - wday) + (n - 1) + elif wday + n <= 4: + # shift by n days without leaving the current week + days = n + else: + # shift by n days plus 2 to get past the weekend + days = n + 2 + result = other + timedelta(days=7 * weeks + days) if self.offset: result = result + self.offset return result @@ -614,8 +617,8 @@ class BusinessHourMixin(BusinessMixin): def __init__(self, start='09:00', end='17:00', offset=timedelta(0)): # must be validated here to equality check kwds = {'offset': offset} - self.start = kwds['start'] = _validate_business_time(start) - self.end = kwds['end'] = _validate_business_time(end) + self.start = kwds['start'] = liboffsets._validate_business_time(start) + self.end = kwds['end'] = liboffsets._validate_business_time(end) self.kwds.update(kwds) self._offset = offset @@ -1092,21 +1095,20 @@ class CustomBusinessMonthBegin(_CustomBusinessMonth): @apply_wraps def apply(self, other): n = self.n - dt_in = other # First move to month offset - cur_mbegin = self.m_offset.rollback(dt_in) + cur_mbegin = self.m_offset.rollback(other) # Find this custom month offset cur_cmbegin = self.cbday.rollforward(cur_mbegin) # handle zero case. arbitrarily rollforward - if n == 0 and dt_in != cur_cmbegin: + if n == 0 and other != cur_cmbegin: n += 1 - if dt_in > cur_cmbegin and n <= -1: + if other > cur_cmbegin and n <= -1: n += 1 - elif dt_in < cur_cmbegin and n >= 1: + elif other < cur_cmbegin and n >= 1: n -= 1 new = cur_mbegin + n * self.m_offset @@ -1239,7 +1241,7 @@ def _apply(self, n, other): months = n // 2 day = 31 if n % 2 else self.day_of_month - return shift_month(other, months, day) + return shift_month(other, months, day_opt=day) def _get_roll(self, i, before_day_of_month, after_day_of_month): n = self.n @@ -1290,7 +1292,7 @@ def _apply(self, n, other): months = n // 2 + n % 2 day = 1 if n % 2 else self.day_of_month - return shift_month(other, months, day) + return shift_month(other, months, day_opt=day) def _get_roll(self, i, before_day_of_month, after_day_of_month): n = self.n @@ -1564,7 +1566,8 @@ class QuarterOffset(DateOffset): _from_name_startingMonth = None _adjust_dst = True # TODO: Consider combining QuarterOffset and YearOffset __init__ at some - # point + # point. Also apply_index, onOffset, rule_code if + # startingMonth vs month attr names are resolved def __init__(self, n=1, normalize=False, startingMonth=None): self.n = self._validate_n(n) @@ -1613,8 +1616,8 @@ def apply(self, other): def onOffset(self, dt): if self.normalize and not _is_normalized(dt): return False - modMonth = (dt.month - self.startingMonth) % 3 - return modMonth == 0 and dt.day == self._get_offset_day(dt) + mod_month = (dt.month - self.startingMonth) % 3 + return mod_month == 0 and dt.day == self._get_offset_day(dt) @apply_index_wraps def apply_index(self, dtindex): @@ -2158,6 +2161,7 @@ def apply(self, other): n -= 1 elif n < 0 and other > current_easter: n += 1 + # TODO: Why does this handle the 0 case the opposite of others? # NOTE: easter returns a datetime.date so we have to convert to type of # other