Skip to content

Commit

Permalink
Follow-Up: Unify apply and onOffset implementations (pandas-dev#18329)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored and jreback committed Nov 19, 2017
1 parent b00e62c commit c868423
Show file tree
Hide file tree
Showing 3 changed files with 118 additions and 63 deletions.
29 changes: 26 additions & 3 deletions pandas/_libs/tslibs/offsets.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -357,6 +357,7 @@ class _BaseOffset(object):
_typ = "dateoffset"
_normalize_cache = True
_cacheable = False
_day_opt = None

def __call__(self, other):
return self.apply(other)
Expand Down Expand Up @@ -394,6 +395,11 @@ class _BaseOffset(object):
out = '<%s' % n_str + className + plural + self._repr_attrs() + '>'
return out

def _get_offset_day(self, datetime other):
# subclass must implement `_day_opt`; calling from the base class
# will raise NotImplementedError.
return get_day_of_month(other, self._day_opt)


class BaseOffset(_BaseOffset):
# Here we add __rfoo__ methods that don't play well with cdef classes
Expand Down Expand Up @@ -468,7 +474,7 @@ cpdef datetime shift_month(datetime stamp, int months, object day_opt=None):
return stamp.replace(year=year, month=month, day=day)


cdef int get_day_of_month(datetime other, day_opt) except? -1:
cpdef int get_day_of_month(datetime other, day_opt) except? -1:
"""
Find the day in `other`'s month that satisfies a DateOffset's onOffset
policy, as described by the `day_opt` argument.
Expand All @@ -493,10 +499,27 @@ cdef int get_day_of_month(datetime other, day_opt) except? -1:
30
"""
cdef:
int wkday, days_in_month

if day_opt == 'start':
return 1
elif day_opt == 'end':
return monthrange(other.year, other.month)[1]

wkday, days_in_month = monthrange(other.year, other.month)
if day_opt == 'end':
return days_in_month
elif day_opt == 'business_start':
# first business day of month
return get_firstbday(wkday, days_in_month)
elif day_opt == 'business_end':
# last business day of month
return get_lastbday(wkday, days_in_month)
elif is_integer_object(day_opt):
day = 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
raise NotImplementedError
else:
raise ValueError(day_opt)

Expand Down
8 changes: 8 additions & 0 deletions pandas/tests/tseries/offsets/test_offsets.py
Original file line number Diff line number Diff line change
Expand Up @@ -4680,3 +4680,11 @@ 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_offset_day_error():
# subclass of _BaseOffset must override _day_opt attribute, or we should
# get a NotImplementedError

with pytest.raises(NotImplementedError):
DateOffset()._get_offset_day(datetime.now())
Loading

0 comments on commit c868423

Please sign in to comment.