Cleanup & De-duplication for custom offset classes (#18640)

Author: jbrockmendel

pandas/_libs/tslibs/offsets.pyx

@@ -307,27 +307,6 @@ class CacheableOffset(object):
     _cacheable = True
 
 
-class BeginMixin(object):
-    # helper for vectorized offsets
-
-    def _beg_apply_index(self, i, freq):
-        """Offsets index to beginning of Period frequency"""
-
-        off = i.to_perioddelta('D')
-
-        base, mult = get_freq_code(freq)
-        base_period = i.to_period(base)
-        if self.n <= 0:
-            # when subtracting, dates on start roll to prior
-            roll = np.where(base_period.to_timestamp() == i - off,
-                            self.n, self.n + 1)
-        else:
-            roll = self.n
-
-        base = (base_period + roll).to_timestamp()
-        return base + off
-
-
 class EndMixin(object):
     # helper for vectorized offsets
 

pandas/tseries/offsets.py

@@ -413,8 +413,27 @@ def _from_name(cls, suffix=None):
         return cls()
 
 
+class _CustomMixin(object):
+    """
+    Mixin for classes that define and validate calendar, holidays,
+    and weekdays attributes
+    """
+    def __init__(self, weekmask, holidays, calendar):
+        calendar, holidays = _get_calendar(weekmask=weekmask,
+                                           holidays=holidays,
+                                           calendar=calendar)
+        # Custom offset instances are identified by the
+        # following two attributes. See DateOffset._params()
+        # holidays, weekmask
+
+        # assumes self.kwds already exists
+        self.kwds['weekmask'] = self.weekmask = weekmask
+        self.kwds['holidays'] = self.holidays = holidays
+        self.kwds['calendar'] = self.calendar = calendar
+
+
 class BusinessMixin(object):
-    """ mixin to business types to provide related functions """
+    """ Mixin to business types to provide related functions """
 
     @property
     def offset(self):
@@ -572,9 +591,26 @@ def __init__(self, start='09:00', end='17:00', offset=timedelta(0)):
         kwds = {'offset': offset}
         self.start = kwds['start'] = _validate_business_time(start)
         self.end = kwds['end'] = _validate_business_time(end)
-        self.kwds = kwds
+        self.kwds.update(kwds)
         self._offset = offset
 
+    @cache_readonly
+    def next_bday(self):
+        """used for moving to next businessday"""
+        if self.n >= 0:
+            nb_offset = 1
+        else:
+            nb_offset = -1
+        if self._prefix.startswith('C'):
+            # CustomBusinessHour
+            return CustomBusinessDay(n=nb_offset,
+                                     weekmask=self.weekmask,
+                                     holidays=self.holidays,
+                                     calendar=self.calendar)
+        else:
+            return BusinessDay(n=nb_offset)
+
+    # TODO: Cache this once offsets are immutable
     def _get_daytime_flag(self):
         if self.start == self.end:
             raise ValueError('start and end must not be the same')
@@ -616,6 +652,7 @@ def _prev_opening_time(self, other):
         return datetime(other.year, other.month, other.day,
                         self.start.hour, self.start.minute)
 
+    # TODO: cache this once offsets are immutable
     def _get_business_hours_by_sec(self):
         """
         Return business hours in a day by seconds.
@@ -784,19 +821,11 @@ def __init__(self, n=1, normalize=False, start='09:00',
                  end='17:00', offset=timedelta(0)):
         self.n = self._validate_n(n)
         self.normalize = normalize
+        self.kwds = {}
         super(BusinessHour, self).__init__(start=start, end=end, offset=offset)
 
-    @cache_readonly
-    def next_bday(self):
-        # used for moving to next businessday
-        if self.n >= 0:
-            nb_offset = 1
-        else:
-            nb_offset = -1
-        return BusinessDay(n=nb_offset)
-
 
-class CustomBusinessDay(BusinessDay):
+class CustomBusinessDay(_CustomMixin, BusinessDay):
     """
     DateOffset subclass representing possibly n custom business days,
     excluding holidays
@@ -822,19 +851,9 @@ def __init__(self, n=1, normalize=False, weekmask='Mon Tue Wed Thu Fri',
         self.n = self._validate_n(n)
         self.normalize = normalize
         self._offset = offset
-        self.kwds = {}
-
-        calendar, holidays = _get_calendar(weekmask=weekmask,
-                                           holidays=holidays,
-                                           calendar=calendar)
-        # CustomBusinessDay instances are identified by the
-        # following two attributes. See DateOffset._params()
-        # holidays, weekmask
+        self.kwds = {'offset': offset}
 
-        self.kwds['weekmask'] = self.weekmask = weekmask
-        self.kwds['holidays'] = self.holidays = holidays
-        self.kwds['calendar'] = self.calendar = calendar
-        self.kwds['offset'] = offset
+        _CustomMixin.__init__(self, weekmask, holidays, calendar)
 
     @apply_wraps
     def apply(self, other):
@@ -874,7 +893,8 @@ def onOffset(self, dt):
         return np.is_busday(day64, busdaycal=self.calendar)
 
 
-class CustomBusinessHour(BusinessHourMixin, SingleConstructorOffset):
+class CustomBusinessHour(_CustomMixin, BusinessHourMixin,
+                         SingleConstructorOffset):
     """
     DateOffset subclass representing possibly n custom business days
 
@@ -889,27 +909,11 @@ def __init__(self, n=1, normalize=False, weekmask='Mon Tue Wed Thu Fri',
                  start='09:00', end='17:00', offset=timedelta(0)):
         self.n = self._validate_n(n)
         self.normalize = normalize
-        super(CustomBusinessHour, self).__init__(start=start,
-                                                 end=end, offset=offset)
-
-        calendar, holidays = _get_calendar(weekmask=weekmask,
-                                           holidays=holidays,
-                                           calendar=calendar)
-        self.kwds['weekmask'] = self.weekmask = weekmask
-        self.kwds['holidays'] = self.holidays = holidays
-        self.kwds['calendar'] = self.calendar = calendar
+        self._offset = offset
+        self.kwds = {'offset': offset}
 
-    @cache_readonly
-    def next_bday(self):
-        # used for moving to next businessday
-        if self.n >= 0:
-            nb_offset = 1
-        else:
-            nb_offset = -1
-        return CustomBusinessDay(n=nb_offset,
-                                 weekmask=self.weekmask,
-                                 holidays=self.holidays,
-                                 calendar=self.calendar)
+        _CustomMixin.__init__(self, weekmask, holidays, calendar)
+        BusinessHourMixin.__init__(self, start=start, end=end, offset=offset)
 
 
 # ---------------------------------------------------------------------
@@ -981,10 +985,10 @@ class BusinessMonthBegin(MonthOffset):
     _day_opt = 'business_start'
 
 
-class CustomBusinessMonthEnd(BusinessMixin, MonthOffset):
+class _CustomBusinessMonth(_CustomMixin, BusinessMixin, MonthOffset):
     """
     DateOffset subclass representing one custom business month, incrementing
-    between end of month dates
+    between [BEGIN/END] of month dates
 
     Parameters
     ----------
@@ -999,27 +1003,19 @@ class CustomBusinessMonthEnd(BusinessMixin, MonthOffset):
         passed to ``numpy.busdaycalendar``
     calendar : pd.HolidayCalendar or np.busdaycalendar
     """
-
     _cacheable = False
-    _prefix = 'CBM'
 
-    onOffset = DateOffset.onOffset  # override MonthOffset method
+    onOffset = DateOffset.onOffset        # override MonthOffset method
     apply_index = DateOffset.apply_index  # override MonthOffset method
 
     def __init__(self, n=1, normalize=False, weekmask='Mon Tue Wed Thu Fri',
                  holidays=None, calendar=None, offset=timedelta(0)):
         self.n = self._validate_n(n)
         self.normalize = normalize
         self._offset = offset
-        self.kwds = {}
+        self.kwds = {'offset': offset}
 
-        calendar, holidays = _get_calendar(weekmask=weekmask,
-                                           holidays=holidays,
-                                           calendar=calendar)
-        self.kwds['weekmask'] = self.weekmask = weekmask
-        self.kwds['holidays'] = self.holidays = holidays
-        self.kwds['calendar'] = self.calendar = calendar
-        self.kwds['offset'] = offset
+        _CustomMixin.__init__(self, weekmask, holidays, calendar)
 
     @cache_readonly
     def cbday(self):
@@ -1028,7 +1024,17 @@ def cbday(self):
 
     @cache_readonly
     def m_offset(self):
-        return MonthEnd(n=1, normalize=self.normalize)
+        if self._prefix.endswith('S'):
+            # MonthBegin:
+            return MonthBegin(n=1, normalize=self.normalize)
+        else:
+            # MonthEnd
+            return MonthEnd(n=1, normalize=self.normalize)
+
+
+class CustomBusinessMonthEnd(_CustomBusinessMonth):
+    __doc__ = _CustomBusinessMonth.__doc__.replace('[BEGIN/END]', 'end')
+    _prefix = 'CBM'
 
     @apply_wraps
     def apply(self, other):
@@ -1054,57 +1060,10 @@ def apply(self, other):
         return result
 
 
-class CustomBusinessMonthBegin(BusinessMixin, MonthOffset):
-    """
-    DateOffset subclass representing one custom business month, incrementing
-    between beginning of month dates
-
-    Parameters
-    ----------
-    n : int, default 1
-    offset : timedelta, default timedelta(0)
-    normalize : bool, default False
-        Normalize start/end dates to midnight before generating date range
-    weekmask : str, Default 'Mon Tue Wed Thu Fri'
-        weekmask of valid business days, passed to ``numpy.busdaycalendar``
-    holidays : list
-        list/array of dates to exclude from the set of valid business days,
-        passed to ``numpy.busdaycalendar``
-    calendar : pd.HolidayCalendar or np.busdaycalendar
-    """
-
-    _cacheable = False
+class CustomBusinessMonthBegin(_CustomBusinessMonth):
+    __doc__ = _CustomBusinessMonth.__doc__.replace('[BEGIN/END]', 'beginning')
     _prefix = 'CBMS'
 
-    onOffset = DateOffset.onOffset  # override MonthOffset method
-    apply_index = DateOffset.apply_index  # override MonthOffset method
-
-    def __init__(self, n=1, normalize=False, weekmask='Mon Tue Wed Thu Fri',
-                 holidays=None, calendar=None, offset=timedelta(0)):
-        self.n = self._validate_n(n)
-        self.normalize = normalize
-        self._offset = offset
-        self.kwds = {}
-
-        # _get_calendar does validation and possible transformation
-        # of calendar and holidays.
-        calendar, holidays = _get_calendar(weekmask=weekmask,
-                                           holidays=holidays,
-                                           calendar=calendar)
-        self.kwds['calendar'] = self.calendar = calendar
-        self.kwds['weekmask'] = self.weekmask = weekmask
-        self.kwds['holidays'] = self.holidays = holidays
-        self.kwds['offset'] = offset
-
-    @cache_readonly
-    def cbday(self):
-        kwds = self.kwds
-        return CustomBusinessDay(n=self.n, normalize=self.normalize, **kwds)
-
-    @cache_readonly
-    def m_offset(self):
-        return MonthBegin(n=1, normalize=self.normalize)
-
     @apply_wraps
     def apply(self, other):
         n = self.n
@@ -1707,13 +1666,16 @@ def onOffset(self, dt):
         return dt.month == self.month and dt.day == self._get_offset_day(dt)
 
     def __init__(self, n=1, normalize=False, month=None):
+        self.n = self._validate_n(n)
+        self.normalize = normalize
+
         month = month if month is not None else self._default_month
         self.month = month
 
         if self.month < 1 or self.month > 12:
             raise ValueError('Month must go from 1 to 12')
 
-        DateOffset.__init__(self, n=n, normalize=normalize, month=month)
+        self.kwds = {'month': month}
 
     @classmethod
     def _from_name(cls, suffix=None):