simplify+unify offset.apply logic (#18263)

Author: jbrockmendel

pandas/_libs/tslibs/offsets.pyx

@@ -139,12 +139,45 @@ def apply_index_wraps(func):
 # ---------------------------------------------------------------------
 # Business Helpers
 
-cpdef int _get_firstbday(int wkday):
+cpdef int get_lastbday(int wkday, int days_in_month):
     """
-    wkday is the result of monthrange(year, month)
+    Find the last day of the month that is a business day.
 
-    If it's a saturday or sunday, increment first business day to reflect this
+    (wkday, days_in_month) is the output from monthrange(year, month)
+
+    Parameters
+    ----------
+    wkday : int
+    days_in_month : int
+
+    Returns
+    -------
+    last_bday : int
     """
+    return days_in_month - max(((wkday + days_in_month - 1) % 7) - 4, 0)
+
+
+cpdef int get_firstbday(int wkday, int days_in_month=0):
+    """
+    Find the first day of the month that is a business day.
+
+    (wkday, days_in_month) is the output from monthrange(year, month)
+
+    Parameters
+    ----------
+    wkday : int
+    days_in_month : int, default 0
+
+    Returns
+    -------
+    first_bday : int
+
+    Notes
+    -----
+    `days_in_month` arg is a dummy so that this has the same signature as
+    `get_lastbday`.
+    """
+    cdef int first
     first = 1
     if wkday == 5:  # on Saturday
         first = 3
@@ -380,7 +413,6 @@ class BaseOffset(_BaseOffset):
 # ----------------------------------------------------------------------
 # RelativeDelta Arithmetic
 
-
 cpdef datetime shift_month(datetime stamp, int months, object day_opt=None):
     """
     Given a datetime (or Timestamp) `stamp`, an integer `months` and an
@@ -406,7 +438,7 @@ cpdef datetime shift_month(datetime stamp, int months, object day_opt=None):
     """
     cdef:
         int year, month, day
-        int dim, dy
+        int wkday, days_in_month, dy
 
     dy = (stamp.month + months) // 12
     month = (stamp.month + months) % 12
@@ -416,15 +448,21 @@ cpdef datetime shift_month(datetime stamp, int months, object day_opt=None):
         dy -= 1
     year = stamp.year + dy
 
-    dim = monthrange(year, month)[1]
+    wkday, days_in_month = monthrange(year, month)
     if day_opt is None:
-        day = min(stamp.day, dim)
+        day = min(stamp.day, days_in_month)
     elif day_opt == 'start':
         day = 1
     elif day_opt == 'end':
-        day = dim
+        day = days_in_month
+    elif day_opt == 'business_start':
+        # first business day of month
+        day = get_firstbday(wkday, days_in_month)
+    elif day_opt == 'business_end':
+        # last business day of month
+        day = get_lastbday(wkday, days_in_month)
     elif is_integer_object(day_opt):
-        day = min(day_opt, dim)
+        day = min(day_opt, days_in_month)
     else:
         raise ValueError(day_opt)
     return stamp.replace(year=year, month=month, day=day)

pandas/tests/tseries/offsets/test_offsets.py

@@ -34,7 +34,7 @@
     to_datetime, DateParseError)
 import pandas.tseries.offsets as offsets
 from pandas.io.pickle import read_pickle
-from pandas._libs.tslibs import timezones
+from pandas._libs.tslibs import timezones, offsets as liboffsets
 from pandas._libs.tslib import normalize_date, NaT, Timestamp
 import pandas._libs.tslib as tslib
 import pandas.util.testing as tm
@@ -4683,3 +4683,43 @@ def test_all_offset_classes(self, tup):
         first = Timestamp(test_values[0], tz='US/Eastern') + offset()
         second = Timestamp(test_values[1], tz='US/Eastern')
         assert first == second
+
+
+def test_get_lastbday():
+    dt = datetime(2017, 11, 30)
+    assert dt.weekday() == 3  # i.e. this is a business day
+    wkday, days_in_month = tslib.monthrange(dt.year, dt.month)
+    assert liboffsets.get_lastbday(wkday, days_in_month) == 30
+
+    dt = datetime(1993, 10, 31)
+    assert dt.weekday() == 6  # i.e. this is not a business day
+    wkday, days_in_month = tslib.monthrange(dt.year, dt.month)
+    assert liboffsets.get_lastbday(wkday, days_in_month) == 29
+
+
+def test_get_firstbday():
+    dt = datetime(2017, 4, 1)
+    assert dt.weekday() == 5  # i.e. not a weekday
+    wkday, days_in_month = tslib.monthrange(dt.year, dt.month)
+    assert liboffsets.get_firstbday(wkday, days_in_month) == 3
+
+    dt = datetime(1993, 10, 1)
+    assert dt.weekday() == 4  # i.e. a business day
+    wkday, days_in_month = tslib.monthrange(dt.year, dt.month)
+    assert liboffsets.get_firstbday(wkday, days_in_month) == 1
+
+
+def test_shift_month():
+    dt = datetime(2017, 11, 30)
+    assert liboffsets.shift_month(dt, 0, 'business_end') == dt
+    assert liboffsets.shift_month(dt, 0,
+                                  'business_start') == datetime(2017, 11, 1)
+
+    ts = Timestamp('1929-05-05')
+    assert liboffsets.shift_month(ts, 1, 'start') == Timestamp('1929-06-01')
+    assert liboffsets.shift_month(ts, -3, 'end') == Timestamp('1929-02-28')
+
+    assert liboffsets.shift_month(ts, 25, None) == Timestamp('1931-06-5')
+
+    # Try to shift to April 31, then shift back to Apr 30 to get a real date
+    assert liboffsets.shift_month(ts, -1, 31) == Timestamp('1929-04-30')

pandas/tseries/offsets.py

@@ -18,7 +18,8 @@
 from pandas._libs.tslibs.offsets import (
     ApplyTypeError,
     as_datetime, _is_normalized,
-    _get_firstbday, _get_calendar, _to_dt64, _validate_business_time,
+    get_firstbday, get_lastbday,
+    _get_calendar, _to_dt64, _validate_business_time,
     _int_to_weekday, _weekday_to_int,
     _determine_offset,
     apply_index_wraps,
@@ -1181,18 +1182,14 @@ class BusinessMonthEnd(MonthOffset):
     def apply(self, other):
         n = self.n
         wkday, days_in_month = tslib.monthrange(other.year, other.month)
-        lastBDay = days_in_month - max(((wkday + days_in_month - 1)
-                                        % 7) - 4, 0)
+        lastBDay = get_lastbday(wkday, days_in_month)
 
         if n > 0 and not other.day >= lastBDay:
             n = n - 1
         elif n <= 0 and other.day > lastBDay:
             n = n + 1
-        other = shift_month(other, n, 'end')
 
-        if other.weekday() > 4:
-            other = other - BDay()
-        return other
+        return shift_month(other, n, 'business_end')
 
 
 class BusinessMonthBegin(MonthOffset):
@@ -1203,7 +1200,7 @@ class BusinessMonthBegin(MonthOffset):
     def apply(self, other):
         n = self.n
         wkday, _ = tslib.monthrange(other.year, other.month)
-        first = _get_firstbday(wkday)
+        first = get_firstbday(wkday)
 
         if other.day > first and n <= 0:
             # as if rolled forward already
@@ -1212,24 +1209,13 @@ def apply(self, other):
             other = other + timedelta(days=first - other.day)
             n -= 1
 
-        other = shift_month(other, n, None)
-        wkday, _ = tslib.monthrange(other.year, other.month)
-        first = _get_firstbday(wkday)
-        result = datetime(other.year, other.month, first,
-                          other.hour, other.minute,
-                          other.second, other.microsecond)
-        return result
+        return shift_month(other, n, 'business_start')
 
     def onOffset(self, dt):
         if self.normalize and not _is_normalized(dt):
             return False
         first_weekday, _ = tslib.monthrange(dt.year, dt.month)
-        if first_weekday == 5:
-            return dt.day == 3
-        elif first_weekday == 6:
-            return dt.day == 2
-        else:
-            return dt.day == 1
+        return dt.day == get_firstbday(first_weekday)
 
 
 class CustomBusinessMonthEnd(BusinessMixin, MonthOffset):
@@ -1611,10 +1597,7 @@ def _from_name(cls, suffix=None):
 
 class QuarterOffset(DateOffset):
     """Quarter representation - doesn't call super"""
-
-    #: default month for __init__
     _default_startingMonth = None
-    #: default month in _from_name
     _from_name_startingMonth = None
     _adjust_dst = True
     # TODO: Consider combining QuarterOffset and YearOffset __init__ at some
@@ -1656,21 +1639,15 @@ class BQuarterEnd(QuarterOffset):
     """
     _outputName = 'BusinessQuarterEnd'
     _default_startingMonth = 3
-    # 'BQ'
     _from_name_startingMonth = 12
     _prefix = 'BQ'
 
     @apply_wraps
     def apply(self, other):
         n = self.n
-        base = other
-        other = datetime(other.year, other.month, other.day,
-                         other.hour, other.minute, other.second,
-                         other.microsecond)
 
         wkday, days_in_month = tslib.monthrange(other.year, other.month)
-        lastBDay = days_in_month - max(((wkday + days_in_month - 1)
-                                        % 7) - 4, 0)
+        lastBDay = get_lastbday(wkday, days_in_month)
 
         monthsToGo = 3 - ((other.month - self.startingMonth) % 3)
         if monthsToGo == 3:
@@ -1681,11 +1658,7 @@ def apply(self, other):
         elif n <= 0 and other.day > lastBDay and monthsToGo == 0:
             n = n + 1
 
-        other = shift_month(other, monthsToGo + 3 * n, 'end')
-        other = tslib._localize_pydatetime(other, base.tzinfo)
-        if other.weekday() > 4:
-            other = other - BDay()
-        return other
+        return shift_month(other, monthsToGo + 3 * n, 'business_end')
 
     def onOffset(self, dt):
         if self.normalize and not _is_normalized(dt):
@@ -1711,7 +1684,7 @@ def apply(self, other):
         n = self.n
         wkday, _ = tslib.monthrange(other.year, other.month)
 
-        first = _get_firstbday(wkday)
+        first = get_firstbday(wkday)
 
         monthsSince = (other.month - self.startingMonth) % 3
 
@@ -1725,14 +1698,7 @@ def apply(self, other):
         elif n > 0 and (monthsSince == 0 and other.day < first):
             n = n - 1
 
-        # get the first bday for result
-        other = shift_month(other, 3 * n - monthsSince, None)
-        wkday, _ = tslib.monthrange(other.year, other.month)
-        first = _get_firstbday(wkday)
-        result = datetime(other.year, other.month, first,
-                          other.hour, other.minute, other.second,
-                          other.microsecond)
-        return result
+        return shift_month(other, 3 * n - monthsSince, 'business_start')
 
 
 class QuarterEnd(EndMixin, QuarterOffset):
@@ -1841,8 +1807,7 @@ class BYearEnd(YearOffset):
     def apply(self, other):
         n = self.n
         wkday, days_in_month = tslib.monthrange(other.year, self.month)
-        lastBDay = (days_in_month -
-                    max(((wkday + days_in_month - 1) % 7) - 4, 0))
+        lastBDay = get_lastbday(wkday, days_in_month)
 
         years = n
         if n > 0:
@@ -1854,17 +1819,8 @@ def apply(self, other):
                     (other.month == self.month and other.day > lastBDay)):
                 years += 1
 
-        other = shift_month(other, 12 * years, None)
-
-        _, days_in_month = tslib.monthrange(other.year, self.month)
-        result = datetime(other.year, self.month, days_in_month,
-                          other.hour, other.minute, other.second,
-                          other.microsecond)
-
-        if result.weekday() > 4:
-            result = result - BDay()
-
-        return result
+        months = years * 12 + (self.month - other.month)
+        return shift_month(other, months, 'business_end')
 
 
 class BYearBegin(YearOffset):
@@ -1878,7 +1834,7 @@ def apply(self, other):
         n = self.n
         wkday, days_in_month = tslib.monthrange(other.year, self.month)
 
-        first = _get_firstbday(wkday)
+        first = get_firstbday(wkday)
 
         years = n
 
@@ -1892,11 +1848,8 @@ def apply(self, other):
                 years += 1
 
         # set first bday for result
-        other = shift_month(other, years * 12, None)
-        wkday, days_in_month = tslib.monthrange(other.year, self.month)
-        first = _get_firstbday(wkday)
-        return datetime(other.year, self.month, first, other.hour,
-                        other.minute, other.second, other.microsecond)
+        months = years * 12 + (self.month - other.month)
+        return shift_month(other, months, 'business_start')
 
 
 class YearEnd(EndMixin, YearOffset):