diff --git a/pandas/_libs/tslibs/offsets.pyx b/pandas/_libs/tslibs/offsets.pyx index 2d8ce4c59fedc..6071241f98b9b 100644 --- a/pandas/_libs/tslibs/offsets.pyx +++ b/pandas/_libs/tslibs/offsets.pyx @@ -139,12 +139,45 @@ 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) + Find the last day of the month that is a business day. - If it's a saturday or sunday, increment first business day to reflect this + (wkday, days_in_month) is the output from monthrange(year, month) + + Parameters + ---------- + wkday : int + days_in_month : int + + Returns + ------- + last_bday : int """ + return days_in_month - max(((wkday + days_in_month - 1) % 7) - 4, 0) + + +cpdef int get_firstbday(int wkday, int days_in_month=0): + """ + 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 + + 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 first = 1 if wkday == 5: # on Saturday first = 3 @@ -380,7 +413,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 +438,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 +448,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) diff --git a/pandas/tests/tseries/offsets/test_offsets.py b/pandas/tests/tseries/offsets/test_offsets.py index b123fa127e29c..250e57c273603 100644 --- a/pandas/tests/tseries/offsets/test_offsets.py +++ b/pandas/tests/tseries/offsets/test_offsets.py @@ -34,7 +34,7 @@ to_datetime, DateParseError) import pandas.tseries.offsets as offsets from pandas.io.pickle import read_pickle -from pandas._libs.tslibs import timezones +from pandas._libs.tslibs import timezones, offsets as liboffsets from pandas._libs.tslib import normalize_date, NaT, Timestamp import pandas._libs.tslib as tslib import pandas.util.testing as tm @@ -4683,3 +4683,43 @@ def test_all_offset_classes(self, tup): first = Timestamp(test_values[0], tz='US/Eastern') + offset() second = Timestamp(test_values[1], tz='US/Eastern') assert first == second + + +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 + + 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 + + +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 + + 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 + + +def test_shift_month(): + dt = datetime(2017, 11, 30) + assert liboffsets.shift_month(dt, 0, 'business_end') == dt + assert liboffsets.shift_month(dt, 0, + 'business_start') == datetime(2017, 11, 1) + + ts = Timestamp('1929-05-05') + assert liboffsets.shift_month(ts, 1, 'start') == Timestamp('1929-06-01') + assert liboffsets.shift_month(ts, -3, 'end') == Timestamp('1929-02-28') + + assert liboffsets.shift_month(ts, 25, None) == Timestamp('1931-06-5') + + # Try to shift to April 31, then shift back to Apr 30 to get a real date + assert liboffsets.shift_month(ts, -1, 31) == Timestamp('1929-04-30') diff --git a/pandas/tseries/offsets.py b/pandas/tseries/offsets.py index 5f94c0cf5a638..ff6eb0d11a6f7 100644 --- a/pandas/tseries/offsets.py +++ b/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, @@ -1180,18 +1181,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): @@ -1202,7 +1199,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 @@ -1211,24 +1208,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): @@ -1610,10 +1596,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 @@ -1655,21 +1638,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: @@ -1680,11 +1657,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): @@ -1710,7 +1683,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 @@ -1724,14 +1697,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): @@ -1840,8 +1806,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: @@ -1853,17 +1818,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): @@ -1877,7 +1833,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 @@ -1891,11 +1847,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):