implement business_start, business_end for shift_months

Author: jbrockmendel

pandas/_libs/tslibs/ccalendar.pxd

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+# cython: profile=False
+
+from cython cimport Py_ssize_t
+
+from numpy cimport int64_t, int32_t
+
+
+cpdef monthrange(int64_t year, Py_ssize_t month)
+
+cdef int dayofweek(int y, int m, int d) nogil
+cdef int is_leapyear(int64_t year) nogil
+cdef int32_t get_days_in_month(int year, int month) nogil

pandas/_libs/tslibs/ccalendar.pyx

@@ -0,0 +1,65 @@
+# -*- coding: utf-8 -*-
+# cython: profile=False
+# cython: boundscheck=False
+"""
+Cython implementations of functions resembling the stdlib calendar module
+"""
+
+cimport cython
+from cython cimport Py_ssize_t
+
+import numpy as np
+cimport numpy as np
+from numpy cimport int64_t, int32_t
+np.import_array()
+
+
+# ----------------------------------------------------------------------
+# Constants
+
+# Slightly more performant cython lookups than a 2D table
+cdef int32_t* days_per_month_array = [
+    31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31, 
+    31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
+
+# ----------------------------------------------------------------------
+
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+cdef inline int32_t get_days_in_month(int year, int month) nogil:
+    return days_per_month_array[12 * is_leapyear(year) + month - 1]
+
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+cpdef monthrange(int64_t year, Py_ssize_t month):
+    cdef:
+        int32_t days
+
+    if month < 1 or month > 12:
+        raise ValueError("bad month number 0; must be 1-12")
+
+    days = get_days_in_month(year, month)
+    return (dayofweek(year, month, 1), days)
+
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+@cython.cdivision
+cdef int dayofweek(int y, int m, int d) nogil:
+    """Sakamoto's method, from wikipedia"""
+    cdef:
+        int day
+        int* sakamoto_arr = [0, 3, 2, 5, 0, 3, 5, 1, 4, 6, 2, 4]
+
+    y -= m < 3
+    day = (y + y / 4 - y / 100 + y / 400 + sakamoto_arr[m - 1] + d) % 7
+    # convert to python day
+    return (day + 6) % 7
+
+
+cdef int is_leapyear(int64_t year) nogil:
+    """Returns 1 if the given year is a leap year, 0 otherwise."""
+    return ((year & 0x3) == 0 and  # year % 4 == 0 
+            ((year % 100) != 0 or (year % 400) == 0))

pandas/_libs/tslibs/offsets.pyx

@@ -17,14 +17,12 @@ np.import_array()
 
 from util cimport is_string_object, is_integer_object
 
-from pandas._libs.tslib import monthrange
-
+from ccalendar cimport get_days_in_month, monthrange
 from conversion cimport tz_convert_single, pydt_to_i8
 from frequencies cimport get_freq_code
 from nattype cimport NPY_NAT
 from np_datetime cimport (pandas_datetimestruct,
-                          dtstruct_to_dt64, dt64_to_dtstruct,
-                          is_leapyear, days_per_month_table)
+                          dtstruct_to_dt64, dt64_to_dtstruct)
 
 # ---------------------------------------------------------------------
 # Constants
@@ -425,11 +423,6 @@ class BaseOffset(_BaseOffset):
 # ----------------------------------------------------------------------
 # RelativeDelta Arithmetic
 
-@cython.wraparound(False)
-@cython.boundscheck(False)
-cdef inline int get_days_in_month(int year, int month) nogil:
-    return days_per_month_table[is_leapyear(year)][month - 1]
-
 
 cdef inline int year_add_months(pandas_datetimestruct dts, int months) nogil:
     """new year number after shifting pandas_datetimestruct number of months"""
@@ -527,8 +520,58 @@ def shift_months(int64_t[:] dtindex, int months, object day=None):
 
                 dts.day = get_days_in_month(dts.year, dts.month)
                 out[i] = dtstruct_to_dt64(&dts)
+
+    elif day == 'business_start':
+        for i in range(count):
+            if dtindex[i] == NPY_NAT:
+                out[i] = NPY_NAT
+                continue
+
+            dt64_to_dtstruct(dtindex[i], &dts)
+            months_to_roll = months
+            wkday, days_in_month = monthrange(dts.year, dts.month)
+            compare_day = get_firstbday(wkday, days_in_month)
+
+            if months_to_roll > 0 and dts.day < compare_day:
+                months_to_roll -= 1
+            elif months_to_roll <= 0 and dts.day > compare_day:
+                # as if rolled forward already
+                months_to_roll += 1
+
+            dts.year = year_add_months(dts, months_to_roll)
+            dts.month = month_add_months(dts, months_to_roll)
+
+            wkday, days_in_month = monthrange(dts.year, dts.month)
+            dts.day = get_firstbday(wkday, days_in_month)
+            out[i] = dtstruct_to_dt64(&dts)
+
+    elif day == 'business_end':
+        for i in range(count):
+            if dtindex[i] == NPY_NAT:
+                out[i] = NPY_NAT
+                continue
+
+            dt64_to_dtstruct(dtindex[i], &dts)
+            months_to_roll = months
+            wkday, days_in_month = monthrange(dts.year, dts.month)
+            compare_day = get_lastbday(wkday, days_in_month)
+
+            if months_to_roll > 0 and dts.day < compare_day:
+                months_to_roll -= 1
+            elif months_to_roll <= 0 and dts.day > compare_day:
+                # as if rolled forward already
+                months_to_roll += 1
+
+            dts.year = year_add_months(dts, months_to_roll)
+            dts.month = month_add_months(dts, months_to_roll)
+
+            wkday, days_in_month = monthrange(dts.year, dts.month)
+            dts.day = get_lastbday(wkday, days_in_month)
+            out[i] = dtstruct_to_dt64(&dts)
+
     else:
-        raise ValueError("day must be None, 'start' or 'end'")
+        raise ValueError("day must be None, 'start', 'end', "
+                         "'business_start', or 'business_end'")
 
     return np.asarray(out)
 

pandas/tseries/offsets.py

@@ -924,8 +924,9 @@ def name(self):
         if self.isAnchored:
             return self.rule_code
         else:
+            month = liboffsets._int_to_month[self.n]
             return "{code}-{month}".format(code=self.rule_code,
-                                           month=_int_to_month[self.n])
+                                           month=month)
 
     def onOffset(self, dt):
         if self.normalize and not _is_normalized(dt):
@@ -945,28 +946,23 @@ def apply(self, other):
 
         return shift_month(other, n, self._day_opt)
 
+    @apply_index_wraps
+    def apply_index(self, i):
+        shifted = liboffsets.shift_months(i.asi8, self.n, self._day_opt)
+        return i._shallow_copy(shifted)
+
 
 class MonthEnd(MonthOffset):
     """DateOffset of one month end"""
     _prefix = 'M'
     _day_opt = 'end'
 
-    @apply_index_wraps
-    def apply_index(self, i):
-        shifted = liboffsets.shift_months(i.asi8, self.n, self._day_opt)
-        return i._shallow_copy(shifted)
-
 
 class MonthBegin(MonthOffset):
     """DateOffset of one month at beginning"""
     _prefix = 'MS'
     _day_opt = 'start'
 
-    @apply_index_wraps
-    def apply_index(self, i):
-        shifted = liboffsets.shift_months(i.asi8, self.n, self._day_opt)
-        return i._shallow_copy(shifted)
-
 
 class BusinessMonthEnd(MonthOffset):
     """DateOffset increments between business EOM dates"""
@@ -1202,6 +1198,7 @@ class CustomBusinessMonthEnd(BusinessMixin, MonthOffset):
     _prefix = 'CBM'
 
     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)):
@@ -1275,6 +1272,7 @@ class CustomBusinessMonthBegin(BusinessMixin, MonthOffset):
     _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)):
@@ -1590,15 +1588,15 @@ def isAnchored(self):
     def _from_name(cls, suffix=None):
         kwargs = {}
         if suffix:
-            kwargs['startingMonth'] = _month_to_int[suffix]
+            kwargs['startingMonth'] = liboffsets._month_to_int[suffix]
         else:
             if cls._from_name_startingMonth is not None:
                 kwargs['startingMonth'] = cls._from_name_startingMonth
         return cls(**kwargs)
 
     @property
     def rule_code(self):
-        month = _int_to_month[self.startingMonth]
+        month = liboffsets._int_to_month[self.startingMonth]
         return '{prefix}-{month}'.format(prefix=self._prefix, month=month)
 
     @apply_wraps
@@ -1618,6 +1616,12 @@ def apply(self, other):
 
         return shift_month(other, 3 * n - months_since, self._day_opt)
 
+    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)
+
 
 class BQuarterEnd(QuarterOffset):
     """DateOffset increments between business Quarter dates
@@ -1631,16 +1635,6 @@ class BQuarterEnd(QuarterOffset):
     _prefix = 'BQ'
     _day_opt = 'business_end'
 
-    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)
-
-
-_int_to_month = tslib._MONTH_ALIASES
-_month_to_int = {v: k for k, v in _int_to_month.items()}
-
 
 # TODO: This is basically the same as BQuarterEnd
 class BQuarterBegin(QuarterOffset):
@@ -1667,12 +1661,6 @@ class QuarterEnd(EndMixin, QuarterOffset):
     def apply_index(self, i):
         return self._end_apply_index(i, self.freqstr)
 
-    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)
-
 
 class QuarterBegin(BeginMixin, QuarterOffset):
     _outputName = 'QuarterBegin'
@@ -1684,7 +1672,8 @@ class QuarterBegin(BeginMixin, QuarterOffset):
     @apply_index_wraps
     def apply_index(self, i):
         freq_month = 12 if self.startingMonth == 1 else self.startingMonth - 1
-        freqstr = 'Q-{month}'.format(month=_int_to_month[freq_month])
+        month = liboffsets._int_to_month[freq_month]
+        freqstr = 'Q-{month}'.format(month=month)
         return self._beg_apply_index(i, freqstr)
 
 
@@ -1725,12 +1714,12 @@ def __init__(self, n=1, normalize=False, month=None):
     def _from_name(cls, suffix=None):
         kwargs = {}
         if suffix:
-            kwargs['month'] = _month_to_int[suffix]
+            kwargs['month'] = liboffsets._month_to_int[suffix]
         return cls(**kwargs)
 
     @property
     def rule_code(self):
-        month = _int_to_month[self.month]
+        month = liboffsets._int_to_month[self.month]
         return '{prefix}-{month}'.format(prefix=self._prefix, month=month)
 
 
@@ -1771,7 +1760,8 @@ class YearBegin(BeginMixin, YearOffset):
     @apply_index_wraps
     def apply_index(self, i):
         freq_month = 12 if self.month == 1 else self.month - 1
-        freqstr = 'A-{month}'.format(month=_int_to_month[freq_month])
+        month = liboffsets._int_to_month[freq_month]
+        freqstr = 'A-{month}'.format(month=month)
         return self._beg_apply_index(i, freqstr)
 
 
@@ -1956,7 +1946,7 @@ def _get_suffix_prefix(self):
 
     def get_rule_code_suffix(self):
         prefix = self._get_suffix_prefix()
-        month = _int_to_month[self.startingMonth]
+        month = liboffsets._int_to_month[self.startingMonth]
         weekday = _int_to_weekday[self.weekday]
         return '{prefix}-{month}-{weekday}'.format(prefix=prefix, month=month,
                                                    weekday=weekday)
@@ -1971,7 +1961,7 @@ def _parse_suffix(cls, varion_code, startingMonth_code, weekday_code):
             raise ValueError("Unable to parse varion_code: "
                              "{code}".format(code=varion_code))
 
-        startingMonth = _month_to_int[startingMonth_code]
+        startingMonth = liboffsets._month_to_int[startingMonth_code]
         weekday = _weekday_to_int[weekday_code]
 
         return {"weekday": weekday,

setup.py

@@ -343,6 +343,7 @@ class CheckSDist(sdist_class):
                  'pandas/_libs/window.pyx',
                  'pandas/_libs/sparse.pyx',
                  'pandas/_libs/parsers.pyx',
+                 'pandas/_libs/tslibs/ccalendar.pyx',
                  'pandas/_libs/tslibs/strptime.pyx',
                  'pandas/_libs/tslibs/np_datetime.pyx',
                  'pandas/_libs/tslibs/timedeltas.pyx',
@@ -558,6 +559,8 @@ def pxd(name):
                      '_libs/tslibs/nattype'],
         'depends': tseries_depends,
         'sources': np_datetime_sources},
+    '_libs.tslibs.ccalendar': {
+        'pyxfile': '_libs/tslibs/ccalendar'},
     '_libs.tslibs.conversion': {
         'pyxfile': '_libs/tslibs/conversion',
         'pxdfiles': ['_libs/src/util',
@@ -584,6 +587,7 @@ def pxd(name):
     '_libs.tslibs.offsets': {
         'pyxfile': '_libs/tslibs/offsets',
         'pxdfiles': ['_libs/src/util',
+                     '_libs/tslibs/ccalendar',
                      '_libs/tslibs/conversion',
                      '_libs/tslibs/frequencies',
                      '_libs/tslibs/nattype'],