diff --git a/pandas/_libs/tslibs/offsets.pyx b/pandas/_libs/tslibs/offsets.pyx index 654c51f0ca842..6f5ad2ae45f50 100644 --- a/pandas/_libs/tslibs/offsets.pyx +++ b/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 diff --git a/pandas/tests/tseries/offsets/test_liboffsets.py b/pandas/tests/tseries/offsets/test_liboffsets.py index 321104222936b..8aa32bc600ee6 100644 --- a/pandas/tests/tseries/offsets/test_liboffsets.py +++ b/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():