simplify+unify offset.apply logic

pandas/_libs/tslibs/offsets.pyx

@@ -139,12 +139,23 @@ 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)
+    (wkday, days_in_month) is the result of monthrange(year, month)
+    """
+    return days_in_month - max(((wkday + days_in_month - 1) % 7) - 4, 0)
+
+
+cpdef int get_firstbday(int wkday, int days_in_month=0):
+    """
+    (wkday, days_in_month) is the result of monthrange(year, month)
 
     If it's a saturday or sunday, increment first business day to reflect this
+
+    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 +391,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 +416,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 +426,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/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,
@@ -1188,18 +1189,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):
@@ -1210,7 +1207,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
@@ -1219,24 +1216,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):
@@ -1648,10 +1634,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
@@ -1693,21 +1676,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:
@@ -1718,11 +1695,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):
@@ -1748,7 +1721,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
 
@@ -1762,14 +1735,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):
@@ -1878,8 +1844,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:
@@ -1891,17 +1856,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):
@@ -1915,7 +1871,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
 
@@ -1929,11 +1885,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):