move monthrange inside get_first/last_bday, allows nogil (pandas-dev#18512)

Author: jbrockmendel

pandas/_libs/tslibs/offsets.pyx

@@ -24,7 +24,7 @@ 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)
+                          is_leapyear, days_per_month_table, dayofweek)
 
 # ---------------------------------------------------------------------
 # Constants
@@ -145,45 +145,44 @@ def apply_index_wraps(func):
 # ---------------------------------------------------------------------
 # Business Helpers
 
-cpdef int get_lastbday(int wkday, int days_in_month):
+cpdef int get_lastbday(int year, int month) nogil:
     """
     Find the last 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
+    year : int
+    month : int
 
     Returns
     -------
     last_bday : int
     """
+    cdef:
+        int wkday, days_in_month
+
+    wkday = dayofweek(year, month, 1)
+    days_in_month = get_days_in_month(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):
+cpdef int get_firstbday(int year, int month) nogil:
     """
     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
+    year : int
+    month : int
 
     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
+    cdef:
+        int first, wkday
+
+    wkday = dayofweek(year, month, 1)
     first = 1
     if wkday == 5:  # on Saturday
         first = 3
@@ -556,52 +555,50 @@ def shift_months(int64_t[:] dtindex, int months, object day=None):
                 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
+        with nogil:
+            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)
+                dt64_to_dtstruct(dtindex[i], &dts)
+                months_to_roll = months
+                compare_day = get_firstbday(dts.year, dts.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
+                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)
+                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)
+                dts.day = get_firstbday(dts.year, dts.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
+        with nogil:
+            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)
+                dt64_to_dtstruct(dtindex[i], &dts)
+                months_to_roll = months
+                compare_day = get_lastbday(dts.year, dts.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
+                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)
+                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)
+                dts.day = get_lastbday(dts.year, dts.month)
+                out[i] = dtstruct_to_dt64(&dts)
 
     else:
         raise ValueError("day must be None, 'start', 'end', "
@@ -635,7 +632,7 @@ cpdef datetime shift_month(datetime stamp, int months, object day_opt=None):
     """
     cdef:
         int year, month, day
-        int wkday, days_in_month, dy
+        int days_in_month, dy
 
     dy = (stamp.month + months) // 12
     month = (stamp.month + months) % 12
@@ -645,20 +642,21 @@ cpdef datetime shift_month(datetime stamp, int months, object day_opt=None):
         dy -= 1
     year = stamp.year + dy
 
-    wkday, days_in_month = monthrange(year, month)
     if day_opt is None:
+        days_in_month = get_days_in_month(year, month)
         day = min(stamp.day, days_in_month)
     elif day_opt == 'start':
         day = 1
     elif day_opt == 'end':
-        day = days_in_month
+        day = get_days_in_month(year, month)
     elif day_opt == 'business_start':
         # first business day of month
-        day = get_firstbday(wkday, days_in_month)
+        day = get_firstbday(year, month)
     elif day_opt == 'business_end':
         # last business day of month
-        day = get_lastbday(wkday, days_in_month)
+        day = get_lastbday(year, month)
     elif is_integer_object(day_opt):
+        days_in_month = get_days_in_month(year, month)
         day = min(day_opt, days_in_month)
     else:
         raise ValueError(day_opt)
@@ -691,22 +689,22 @@ cpdef int get_day_of_month(datetime other, day_opt) except? -1:
 
     """
     cdef:
-        int wkday, days_in_month
+        int days_in_month
 
     if day_opt == 'start':
         return 1
-
-    wkday, days_in_month = monthrange(other.year, other.month)
-    if day_opt == 'end':
+    elif day_opt == 'end':
+        days_in_month = get_days_in_month(other.year, other.month)
         return days_in_month
     elif day_opt == 'business_start':
         # first business day of month
-        return get_firstbday(wkday, days_in_month)
+        return get_firstbday(other.year, other.month)
     elif day_opt == 'business_end':
         # last business day of month
-        return get_lastbday(wkday, days_in_month)
+        return get_lastbday(other.year, other.month)
     elif is_integer_object(day_opt):
-        day = min(day_opt, days_in_month)
+        days_in_month = get_days_in_month(other.year, other.month)
+        return min(day_opt, days_in_month)
     elif day_opt is None:
         # Note: unlike `shift_month`, get_day_of_month does not
         # allow day_opt = None

pandas/tests/tseries/offsets/test_liboffsets.py

@@ -6,7 +6,6 @@
 
 import pytest
 
-from pandas._libs import tslib
 from pandas import Timestamp
 
 import pandas._libs.tslibs.offsets as liboffsets
@@ -15,25 +14,21 @@
 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
+    assert liboffsets.get_lastbday(dt.year, dt.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
+    assert liboffsets.get_lastbday(dt.year, dt.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
+    assert liboffsets.get_firstbday(dt.year, dt.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
+    assert liboffsets.get_firstbday(dt.year, dt.month) == 1
 
 
 def test_shift_month():