Skip to content

Follow-Up: Unify apply and onOffset implementations #18329

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Nov 19, 2017
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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in a followup (TODO), can you add a description of these attributes in the doc-string of this class. (e.g. what the options are, what they do, etc).

_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