Skip to content

Commit

Permalink
REF: De-duplicate roll_yearday/roll_qtrday (#34760)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored Jun 14, 2020
1 parent a97bddd commit f984364
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 91 deletions.
92 changes: 7 additions & 85 deletions pandas/_libs/tslibs/offsets.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -1879,7 +1879,7 @@ cdef class YearOffset(SingleConstructorOffset):

@apply_wraps
def apply(self, other):
years = roll_yearday(other, self.n, self.month, self._day_opt)
years = roll_qtrday(other, self.n, self.month, self._day_opt, modby=12)
months = years * 12 + (self.month - other.month)
return shift_month(other, months, self._day_opt)

Expand Down Expand Up @@ -4158,10 +4158,13 @@ def roll_qtrday(other: datetime, n: int, month: int,
npy_datetimestruct dts
pydate_to_dtstruct(other, &dts)

# TODO: Merge this with roll_yearday by setting modby=12 there?
# code de-duplication versus perf hit?
# TODO: with small adjustments this could be used in shift_quarters
months_since = other.month % modby - month % modby

if modby == 12:
# We care about the month-of-year, not month-of-quarter, so skip mod
months_since = other.month - month
else:
months_since = other.month % modby - month % modby

if n > 0:
if months_since < 0 or (months_since == 0 and
Expand All @@ -4177,84 +4180,3 @@ def roll_qtrday(other: datetime, n: int, month: int,
# make sure to roll forward, so negate
n += 1
return n


def roll_yearday(other: datetime, n: int, month: int, day_opt: object) -> int:
"""
Possibly increment or decrement the number of periods to shift
based on rollforward/rollbackward conventions.

Parameters
----------
other : datetime or Timestamp
n : number of periods to increment, before adjusting for rolling
month : reference month giving the first month of the year
day_opt : 'start', 'end', 'business_start', 'business_end', or int
The day of the month to compare against that of `other` when
incrementing or decrementing the number of periods:

'start': 1
'end': last day of the month
'business_start': first business day of the month
'business_end': last business day of the month
int: day in the month indicated by `other`, or the last of day
the month if the value exceeds in that month's number of days.

Returns
-------
n : int number of periods to increment

Notes
-----
* Mirrors `roll_check` in shift_months

Examples
-------
>>> month = 3
>>> day_opt = 'start' # `other` will be compared to March 1
>>> other = datetime(2017, 2, 10) # before March 1
>>> roll_yearday(other, 2, month, day_opt)
1
>>> roll_yearday(other, -7, month, day_opt)
-7
>>>
>>> other = Timestamp('2014-03-15', tz='US/Eastern') # after March 1
>>> roll_yearday(other, 2, month, day_opt)
2
>>> roll_yearday(other, -7, month, day_opt)
-6

>>> month = 6
>>> day_opt = 'end' # `other` will be compared to June 30
>>> other = datetime(1999, 6, 29) # before June 30
>>> roll_yearday(other, 5, month, day_opt)
4
>>> roll_yearday(other, -7, month, day_opt)
-7
>>>
>>> other = Timestamp(2072, 8, 24, 6, 17, 18) # after June 30
>>> roll_yearday(other, 5, month, day_opt)
5
>>> roll_yearday(other, -7, month, day_opt)
-6

"""
cdef:
npy_datetimestruct dts
pydate_to_dtstruct(other, &dts)

# Note: The other.day < ... condition will never hold when day_opt=='start'
# and the other.day > ... condition will never hold when day_opt=='end'.
# At some point these extra checks may need to be optimized away.
# But that point isn't today.
if n > 0:
if other.month < month or (other.month == month and
other.day < get_day_of_month(&dts,
day_opt)):
n -= 1
else:
if other.month > month or (other.month == month and
other.day > get_day_of_month(&dts,
day_opt)):
n += 1
return n
12 changes: 6 additions & 6 deletions pandas/tests/tslibs/test_liboffsets.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,11 +88,11 @@ def test_shift_month_error():
],
)
@pytest.mark.parametrize("n", [2, -7, 0])
def test_roll_yearday(other, expected, n):
def test_roll_qtrday_year(other, expected, n):
month = 3
day_opt = "start" # `other` will be compared to March 1.

assert liboffsets.roll_yearday(other, n, month, day_opt) == expected[n]
assert roll_qtrday(other, n, month, day_opt, modby=12) == expected[n]


@pytest.mark.parametrize(
Expand All @@ -105,22 +105,22 @@ def test_roll_yearday(other, expected, n):
],
)
@pytest.mark.parametrize("n", [5, -7, 0])
def test_roll_yearday2(other, expected, n):
def test_roll_qtrday_year2(other, expected, n):
month = 6
day_opt = "end" # `other` will be compared to June 30.

assert liboffsets.roll_yearday(other, n, month, day_opt) == expected[n]
assert roll_qtrday(other, n, month, day_opt, modby=12) == expected[n]


def test_get_day_of_month_error():
# get_day_of_month is not directly exposed.
# We test it via roll_yearday.
# We test it via roll_qtrday.
dt = datetime(2017, 11, 15)
day_opt = "foo"

with pytest.raises(ValueError, match=day_opt):
# To hit the raising case we need month == dt.month and n > 0.
liboffsets.roll_yearday(dt, n=3, month=11, day_opt=day_opt)
roll_qtrday(dt, n=3, month=11, day_opt=day_opt, modby=12)


@pytest.mark.parametrize(
Expand Down

0 comments on commit f984364

Please sign in to comment.