Skip to content

Commit c24a0d2

Browse files
jbrockmendeljreback
jbrockmendel
authored and
jreback
committed
Implement roll_monthday, simplify SemiMonthOffset (#18762)
1 parent faeac49 commit c24a0d2

File tree

3 files changed

+247
-109
lines changed

3 files changed

+247
-109
lines changed

‎pandas/_libs/tslibs/offsets.pyx

+95-9
Original file line numberDiff line numberDiff line change
@@ -523,11 +523,9 @@ def shift_quarters(int64_t[:] dtindex, int quarters,
523523
n = quarters
524524

525525
months_since = (dts.month - q1start_month) % modby
526-
compare_month = dts.month - months_since
527-
compare_month = compare_month or 12
528526
# compare_day is only relevant for comparison in the case
529527
# where months_since == 0.
530-
compare_day = get_firstbday(dts.year, compare_month)
528+
compare_day = get_firstbday(dts.year, dts.month)
531529

532530
if n <= 0 and (months_since != 0 or
533531
(months_since == 0 and dts.day > compare_day)):
@@ -556,11 +554,9 @@ def shift_quarters(int64_t[:] dtindex, int quarters,
556554
n = quarters
557555

558556
months_since = (dts.month - q1start_month) % modby
559-
compare_month = dts.month - months_since
560-
compare_month = compare_month or 12
561557
# compare_day is only relevant for comparison in the case
562558
# where months_since == 0.
563-
compare_day = get_lastbday(dts.year, compare_month)
559+
compare_day = get_lastbday(dts.year, dts.month)
564560

565561
if n <= 0 and (months_since != 0 or
566562
(months_since == 0 and dts.day > compare_day)):
@@ -827,7 +823,55 @@ cpdef int get_day_of_month(datetime other, day_opt) except? -1:
827823
raise ValueError(day_opt)
828824

829825

830-
cpdef int roll_yearday(other, n, month, day_opt='start') except? -1:
826+
cpdef int roll_convention(int other, int n, int compare):
827+
"""
828+
Possibly increment or decrement the number of periods to shift
829+
based on rollforward/rollbackward conventions.
830+
831+
Parameters
832+
----------
833+
other : int, generally the day component of a datetime
834+
n : number of periods to increment, before adjusting for rolling
835+
compare : int, generally the day component of a datetime, in the same
836+
month as the datetime form which `other` was taken.
837+
838+
Returns
839+
-------
840+
n : int number of periods to increment
841+
"""
842+
if n > 0 and other < compare:
843+
n -= 1
844+
elif n <= 0 and other > compare:
845+
# as if rolled forward already
846+
n += 1
847+
return n
848+
849+
850+
cpdef int roll_monthday(datetime other, int n, datetime compare):
851+
"""
852+
Possibly increment or decrement the number of periods to shift
853+
based on rollforward/rollbackward conventions.
854+
855+
Parameters
856+
----------
857+
other : datetime
858+
n : number of periods to increment, before adjusting for rolling
859+
compare : datetime
860+
861+
Returns
862+
-------
863+
n : int number of periods to increment
864+
"""
865+
if n > 0 and other < compare:
866+
n -= 1
867+
elif n <= 0 and other > compare:
868+
# as if rolled forward already
869+
n += 1
870+
return n
871+
872+
873+
cpdef int roll_qtrday(datetime other, int n, int month, object day_opt,
874+
int modby=3) except? -1:
831875
"""
832876
Possibly increment or decrement the number of periods to shift
833877
based on rollforward/rollbackward conventions.
@@ -836,6 +880,48 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1:
836880
----------
837881
other : datetime or Timestamp
838882
n : number of periods to increment, before adjusting for rolling
883+
month : int reference month giving the first month of the year
884+
day_opt : 'start', 'end', 'business_start', 'business_end'
885+
The convention to use in finding the day in a given month against
886+
which to compare for rollforward/rollbackward decisions.
887+
modby : int 3 for quarters, 12 for years
888+
889+
Returns
890+
-------
891+
n : int number of periods to increment
892+
"""
893+
# TODO: Merge this with roll_yearday by setting modby=12 there?
894+
# code de-duplication versus perf hit?
895+
# TODO: with small adjustments this could be used in shift_quarters
896+
months_since = other.month % modby - month % modby
897+
898+
if n > 0:
899+
if months_since < 0 or (months_since == 0 and
900+
other.day < get_day_of_month(other,
901+
day_opt)):
902+
# pretend to roll back if on same month but
903+
# before compare_day
904+
n -= 1
905+
else:
906+
if months_since > 0 or (months_since == 0 and
907+
other.day > get_day_of_month(other,
908+
day_opt)):
909+
# make sure to roll forward, so negate
910+
n += 1
911+
return n
912+
913+
914+
cpdef int roll_yearday(datetime other, int n, int month,
915+
object day_opt) except? -1:
916+
"""
917+
Possibly increment or decrement the number of periods to shift
918+
based on rollforward/rollbackward conventions.
919+
920+
Parameters
921+
----------
922+
other : datetime or Timestamp
923+
n : number of periods to increment, before adjusting for rolling
924+
month : reference month giving the first month of the year
839925
day_opt : 'start', 'end'
840926
'start': returns 1
841927
'end': returns last day of the month
@@ -846,7 +932,7 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1:
846932
847933
Notes
848934
-----
849-
* Mirrors `roll_check` in tslib.shift_months
935+
* Mirrors `roll_check` in shift_months
850936
851937
Examples
852938
-------
@@ -888,7 +974,7 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1:
888974
other.day < get_day_of_month(other,
889975
day_opt)):
890976
n -= 1
891-
elif n <= 0:
977+
else:
892978
if other.month > month or (other.month == month and
893979
other.day > get_day_of_month(other,
894980
day_opt)):

‎pandas/tests/tseries/offsets/test_liboffsets.py

+91
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from pandas import Timestamp
1010

1111
import pandas._libs.tslibs.offsets as liboffsets
12+
from pandas._libs.tslibs.offsets import roll_qtrday
1213

1314

1415
def test_get_lastbday():
@@ -95,3 +96,93 @@ def test_roll_yearday():
9596
assert liboffsets.roll_yearday(other, 5, month, day_opt) == 5
9697
assert liboffsets.roll_yearday(other, -7, month, day_opt) == -6
9798
assert liboffsets.roll_yearday(other, 0, month, day_opt) == 1
99+
100+
101+
def test_roll_qtrday():
102+
other = Timestamp(2072, 10, 1, 6, 17, 18) # Saturday
103+
for day_opt in ['start', 'end', 'business_start', 'business_end']:
104+
# as long as (other.month % 3) != (month % 3), day_opt is irrelevant
105+
# the `day_opt` doesn't matter.
106+
month = 5 # (other.month % 3) < (month % 3)
107+
assert roll_qtrday(other, 4, month, day_opt, modby=3) == 3
108+
assert roll_qtrday(other, -3, month, day_opt, modby=3) == -3
109+
110+
month = 3 # (other.month % 3) > (month % 3)
111+
assert roll_qtrday(other, 4, month, day_opt, modby=3) == 4
112+
assert roll_qtrday(other, -3, month, day_opt, modby=3) == -2
113+
114+
month = 2
115+
other = datetime(1999, 5, 31) # Monday
116+
# has (other.month % 3) == (month % 3)
117+
118+
n = 2
119+
assert roll_qtrday(other, n, month, 'start', modby=3) == n
120+
assert roll_qtrday(other, n, month, 'end', modby=3) == n
121+
assert roll_qtrday(other, n, month, 'business_start', modby=3) == n
122+
assert roll_qtrday(other, n, month, 'business_end', modby=3) == n
123+
124+
n = -1
125+
assert roll_qtrday(other, n, month, 'start', modby=3) == n + 1
126+
assert roll_qtrday(other, n, month, 'end', modby=3) == n
127+
assert roll_qtrday(other, n, month, 'business_start', modby=3) == n + 1
128+
assert roll_qtrday(other, n, month, 'business_end', modby=3) == n
129+
130+
other = Timestamp(2072, 10, 1, 6, 17, 18) # Saturday
131+
month = 4 # (other.month % 3) == (month % 3)
132+
n = 2
133+
assert roll_qtrday(other, n, month, 'start', modby=3) == n
134+
assert roll_qtrday(other, n, month, 'end', modby=3) == n - 1
135+
assert roll_qtrday(other, n, month, 'business_start', modby=3) == n - 1
136+
assert roll_qtrday(other, n, month, 'business_end', modby=3) == n - 1
137+
138+
n = -1
139+
assert roll_qtrday(other, n, month, 'start', modby=3) == n
140+
assert roll_qtrday(other, n, month, 'end', modby=3) == n
141+
assert roll_qtrday(other, n, month, 'business_start', modby=3) == n
142+
assert roll_qtrday(other, n, month, 'business_end', modby=3) == n
143+
144+
other = Timestamp(2072, 10, 3, 6, 17, 18) # First businessday
145+
month = 4 # (other.month % 3) == (month % 3)
146+
n = 2
147+
assert roll_qtrday(other, n, month, 'start', modby=3) == n
148+
assert roll_qtrday(other, n, month, 'end', modby=3) == n - 1
149+
assert roll_qtrday(other, n, month, 'business_start', modby=3) == n
150+
assert roll_qtrday(other, n, month, 'business_end', modby=3) == n - 1
151+
152+
n = -1
153+
assert roll_qtrday(other, n, month, 'start', modby=3) == n + 1
154+
assert roll_qtrday(other, n, month, 'end', modby=3) == n
155+
assert roll_qtrday(other, n, month, 'business_start', modby=3) == n
156+
assert roll_qtrday(other, n, month, 'business_end', modby=3) == n
157+
158+
159+
def test_roll_monthday():
160+
other = Timestamp('2017-12-29', tz='US/Pacific')
161+
before = Timestamp('2017-12-01', tz='US/Pacific')
162+
after = Timestamp('2017-12-31', tz='US/Pacific')
163+
164+
n = 42
165+
assert liboffsets.roll_monthday(other, n, other) == n
166+
assert liboffsets.roll_monthday(other, n, before) == n
167+
assert liboffsets.roll_monthday(other, n, after) == n - 1
168+
169+
n = -4
170+
assert liboffsets.roll_monthday(other, n, other) == n
171+
assert liboffsets.roll_monthday(other, n, before) == n + 1
172+
assert liboffsets.roll_monthday(other, n, after) == n
173+
174+
175+
def test_roll_convention():
176+
other = 29
177+
before = 1
178+
after = 31
179+
180+
n = 42
181+
assert liboffsets.roll_convention(other, n, other) == n
182+
assert liboffsets.roll_convention(other, n, before) == n
183+
assert liboffsets.roll_convention(other, n, after) == n - 1
184+
185+
n = -4
186+
assert liboffsets.roll_convention(other, n, other) == n
187+
assert liboffsets.roll_convention(other, n, before) == n + 1
188+
assert liboffsets.roll_convention(other, n, after) == n

0 commit comments

Comments
 (0)