implement non-controversial cleanup portions of #18762

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,17 +835,20 @@ 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
     -------
     n : int number of periods to increment
 
     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)):

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)

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