Skip to content

Commit

Permalink
simplify CBMonth.apply to remove roll_monthday (pandas-dev#19146)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored and maxim veksler committed Jan 11, 2018
1 parent bb431e1 commit d40074f
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 67 deletions.
23 changes: 0 additions & 23 deletions pandas/_libs/tslibs/offsets.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -847,29 +847,6 @@ cpdef int roll_convention(int other, int n, int compare):
return n


cpdef int roll_monthday(datetime other, int n, datetime compare):
"""
Possibly increment or decrement the number of periods to shift
based on rollforward/rollbackward conventions.
Parameters
----------
other : datetime
n : number of periods to increment, before adjusting for rolling
compare : datetime
Returns
-------
n : int number of periods to increment
"""
if n > 0 and other < compare:
n -= 1
elif n <= 0 and other > compare:
# as if rolled forward already
n += 1
return n


cpdef int roll_qtrday(datetime other, int n, int month, object day_opt,
int modby=3) except? -1:
"""
Expand Down
16 changes: 0 additions & 16 deletions pandas/tests/tseries/offsets/test_liboffsets.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,22 +156,6 @@ def test_roll_qtrday():
assert roll_qtrday(other, n, month, 'business_end', modby=3) == n


def test_roll_monthday():
other = Timestamp('2017-12-29', tz='US/Pacific')
before = Timestamp('2017-12-01', tz='US/Pacific')
after = Timestamp('2017-12-31', tz='US/Pacific')

n = 42
assert liboffsets.roll_monthday(other, n, other) == n
assert liboffsets.roll_monthday(other, n, before) == n
assert liboffsets.roll_monthday(other, n, after) == n - 1

n = -4
assert liboffsets.roll_monthday(other, n, other) == n
assert liboffsets.roll_monthday(other, n, before) == n + 1
assert liboffsets.roll_monthday(other, n, after) == n


def test_roll_convention():
other = 29
before = 1
Expand Down
63 changes: 35 additions & 28 deletions pandas/tseries/offsets.py
Original file line number Diff line number Diff line change
Expand Up @@ -1039,55 +1039,62 @@ def __init__(self, n=1, normalize=False, weekmask='Mon Tue Wed Thu Fri',
_CustomMixin.__init__(self, weekmask, holidays, calendar)

@cache_readonly
def cbday(self):
kwds = self.kwds
return CustomBusinessDay(n=self.n, normalize=self.normalize, **kwds)
def cbday_roll(self):
"""Define default roll function to be called in apply method"""
cbday = CustomBusinessDay(n=self.n, normalize=False, **self.kwds)

if self._prefix.endswith('S'):
# MonthBegin
roll_func = cbday.rollforward
else:
# MonthEnd
roll_func = cbday.rollback
return roll_func

@cache_readonly
def m_offset(self):
if self._prefix.endswith('S'):
# MonthBegin:
return MonthBegin(n=1, normalize=self.normalize)
# MonthBegin
moff = MonthBegin(n=1, normalize=False)
else:
# MonthEnd
return MonthEnd(n=1, normalize=self.normalize)
moff = MonthEnd(n=1, normalize=False)
return moff


class CustomBusinessMonthEnd(_CustomBusinessMonth):
__doc__ = _CustomBusinessMonth.__doc__.replace('[BEGIN/END]', 'end')
_prefix = 'CBM'
@cache_readonly
def month_roll(self):
"""Define default roll function to be called in apply method"""
if self._prefix.endswith('S'):
# MonthBegin
roll_func = self.m_offset.rollback
else:
# MonthEnd
roll_func = self.m_offset.rollforward
return roll_func

@apply_wraps
def apply(self, other):
# First move to month offset
cur_mend = self.m_offset.rollforward(other)
cur_month_offset_date = self.month_roll(other)

# Find this custom month offset
compare_date = self.cbday.rollback(cur_mend)
n = liboffsets.roll_monthday(other, self.n, compare_date)
compare_date = self.cbday_roll(cur_month_offset_date)
n = liboffsets.roll_convention(other.day, self.n, compare_date.day)

new = cur_mend + n * self.m_offset
result = self.cbday.rollback(new)
new = cur_month_offset_date + n * self.m_offset
result = self.cbday_roll(new)
return result


class CustomBusinessMonthEnd(_CustomBusinessMonth):
__doc__ = _CustomBusinessMonth.__doc__.replace('[BEGIN/END]', 'end')
_prefix = 'CBM'


class CustomBusinessMonthBegin(_CustomBusinessMonth):
__doc__ = _CustomBusinessMonth.__doc__.replace('[BEGIN/END]', 'beginning')
_prefix = 'CBMS'

@apply_wraps
def apply(self, other):
# First move to month offset
cur_mbegin = self.m_offset.rollback(other)

# Find this custom month offset
compare_date = self.cbday.rollforward(cur_mbegin)
n = liboffsets.roll_monthday(other, self.n, compare_date)

new = cur_mbegin + n * self.m_offset
result = self.cbday.rollforward(new)
return result


# ---------------------------------------------------------------------
# Semi-Month Based Offset Classes
Expand Down

0 comments on commit d40074f

Please sign in to comment.