Skip to content

REF: de-duplicate code in liboffsets #34778

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 5 commits into from
Jun 15, 2020
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
286 changes: 79 additions & 207 deletions pandas/_libs/tslibs/offsets.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -3723,136 +3723,14 @@ cdef shift_quarters(
out : ndarray[int64_t]
"""
cdef:
Py_ssize_t i
npy_datetimestruct dts
int count = len(dtindex)
int months_to_roll, months_since, n, compare_day
Py_ssize_t count = len(dtindex)
int64_t[:] out = np.empty(count, dtype="int64")

if day_opt == "start":
with nogil:
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue

dt64_to_dtstruct(dtindex[i], &dts)
n = quarters

months_since = (dts.month - q1start_month) % modby
compare_day = get_day_of_month(&dts, day_opt)

# offset semantics - if on the anchor point and going backwards
# shift to next
if n <= 0 and (months_since != 0 or
(months_since == 0 and dts.day > compare_day)):
# make sure to roll forward, so negate
n += 1
elif n > 0 and (months_since == 0 and dts.day < compare_day):
# pretend to roll back if on same month but
# before compare_day
n -= 1

dts.year = year_add_months(dts, modby * n - months_since)
dts.month = month_add_months(dts, modby * n - months_since)
dts.day = get_day_of_month(&dts, day_opt)

out[i] = dtstruct_to_dt64(&dts)

elif day_opt == "end":
with nogil:
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue

dt64_to_dtstruct(dtindex[i], &dts)
n = quarters

months_since = (dts.month - q1start_month) % modby
compare_day = get_day_of_month(&dts, day_opt)

if n <= 0 and (months_since != 0 or
(months_since == 0 and dts.day > compare_day)):
# make sure to roll forward, so negate
n += 1
elif n > 0 and (months_since == 0 and dts.day < compare_day):
# pretend to roll back if on same month but
# before compare_day
n -= 1

dts.year = year_add_months(dts, modby * n - months_since)
dts.month = month_add_months(dts, modby * n - months_since)
dts.day = get_day_of_month(&dts, day_opt)

out[i] = dtstruct_to_dt64(&dts)

elif day_opt == "business_start":
with nogil:
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue

dt64_to_dtstruct(dtindex[i], &dts)
n = quarters

months_since = (dts.month - q1start_month) % modby
# compare_day is only relevant for comparison in the case
# where months_since == 0.
compare_day = get_day_of_month(&dts, day_opt)

if n <= 0 and (months_since != 0 or
(months_since == 0 and dts.day > compare_day)):
# make sure to roll forward, so negate
n += 1
elif n > 0 and (months_since == 0 and dts.day < compare_day):
# pretend to roll back if on same month but
# before compare_day
n -= 1

dts.year = year_add_months(dts, modby * n - months_since)
dts.month = month_add_months(dts, modby * n - months_since)

dts.day = get_day_of_month(&dts, day_opt)

out[i] = dtstruct_to_dt64(&dts)

elif day_opt == "business_end":
with nogil:
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue

dt64_to_dtstruct(dtindex[i], &dts)
n = quarters

months_since = (dts.month - q1start_month) % modby
# compare_day is only relevant for comparison in the case
# where months_since == 0.
compare_day = get_day_of_month(&dts, day_opt)

if n <= 0 and (months_since != 0 or
(months_since == 0 and dts.day > compare_day)):
# make sure to roll forward, so negate
n += 1
elif n > 0 and (months_since == 0 and dts.day < compare_day):
# pretend to roll back if on same month but
# before compare_day
n -= 1

dts.year = year_add_months(dts, modby * n - months_since)
dts.month = month_add_months(dts, modby * n - months_since)

dts.day = get_day_of_month(&dts, day_opt)

out[i] = dtstruct_to_dt64(&dts)

else:
if day_opt not in ["start", "end", "business_start", "business_end"]:
raise ValueError("day must be None, 'start', 'end', "
"'business_start', or 'business_end'")

_shift_quarters(dtindex, out, count, quarters, q1start_month, day_opt, modby)
return np.asarray(out)


Expand All @@ -3872,7 +3750,6 @@ def shift_months(const int64_t[:] dtindex, int months, object day_opt=None):
Py_ssize_t i
npy_datetimestruct dts
int count = len(dtindex)
int months_to_roll
int64_t[:] out = np.empty(count, dtype="int64")

if day_opt is None:
Expand All @@ -3888,94 +3765,90 @@ def shift_months(const int64_t[:] dtindex, int months, object day_opt=None):

dts.day = min(dts.day, get_days_in_month(dts.year, dts.month))
out[i] = dtstruct_to_dt64(&dts)
elif day_opt == "start":
with nogil:
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue

dt64_to_dtstruct(dtindex[i], &dts)
months_to_roll = months
compare_day = get_day_of_month(&dts, day_opt)
elif day_opt in ["start", "end", "business_start", "business_end"]:
_shift_months(dtindex, out, count, months, day_opt)

# offset semantics - if on the anchor point and going backwards
# shift to next
months_to_roll = roll_convention(dts.day, months_to_roll,
compare_day)

dts.year = year_add_months(dts, months_to_roll)
dts.month = month_add_months(dts, months_to_roll)
dts.day = get_day_of_month(&dts, day_opt)

out[i] = dtstruct_to_dt64(&dts)
elif day_opt == "end":
with nogil:
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue
else:
raise ValueError("day must be None, 'start', 'end', "
"'business_start', or 'business_end'")

dt64_to_dtstruct(dtindex[i], &dts)
months_to_roll = months
compare_day = get_day_of_month(&dts, day_opt)
return np.asarray(out)

# similar semantics - when adding shift forward by one
# month if already at an end of month
months_to_roll = roll_convention(dts.day, months_to_roll,
compare_day)

dts.year = year_add_months(dts, months_to_roll)
dts.month = month_add_months(dts, months_to_roll)
@cython.wraparound(False)
@cython.boundscheck(False)
cdef inline void _shift_months(const int64_t[:] dtindex,
int64_t[:] out,
Py_ssize_t count,
int months,
str day_opt) nogil:
"""See shift_months.__doc__"""
cdef:
Py_ssize_t i
int months_to_roll, compare_day
npy_datetimestruct dts

dts.day = get_day_of_month(&dts, day_opt)
out[i] = dtstruct_to_dt64(&dts)
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue

elif day_opt == "business_start":
with nogil:
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue
dt64_to_dtstruct(dtindex[i], &dts)
months_to_roll = months
compare_day = get_day_of_month(&dts, day_opt)

dt64_to_dtstruct(dtindex[i], &dts)
months_to_roll = months
compare_day = get_day_of_month(&dts, day_opt)
months_to_roll = roll_convention(dts.day, months_to_roll,
compare_day)

months_to_roll = roll_convention(dts.day, months_to_roll,
compare_day)
dts.year = year_add_months(dts, months_to_roll)
dts.month = month_add_months(dts, months_to_roll)
dts.day = get_day_of_month(&dts, day_opt)

dts.year = year_add_months(dts, months_to_roll)
dts.month = month_add_months(dts, months_to_roll)
out[i] = dtstruct_to_dt64(&dts)

dts.day = get_day_of_month(&dts, day_opt)
out[i] = dtstruct_to_dt64(&dts)

elif day_opt == "business_end":
with nogil:
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue
@cython.wraparound(False)
@cython.boundscheck(False)
cdef inline void _shift_quarters(const int64_t[:] dtindex,
int64_t[:] out,
Py_ssize_t count,
int quarters,
int q1start_month,
str day_opt,
int modby) nogil:
"""See shift_quarters.__doc__"""
cdef:
Py_ssize_t i
int months_since, compare_day, n
npy_datetimestruct dts

dt64_to_dtstruct(dtindex[i], &dts)
months_to_roll = months
compare_day = get_day_of_month(&dts, day_opt)
for i in range(count):
if dtindex[i] == NPY_NAT:
out[i] = NPY_NAT
continue

months_to_roll = roll_convention(dts.day, months_to_roll,
compare_day)
dt64_to_dtstruct(dtindex[i], &dts)
n = quarters

dts.year = year_add_months(dts, months_to_roll)
dts.month = month_add_months(dts, months_to_roll)
months_since = (dts.month - q1start_month) % modby
compare_day = get_day_of_month(&dts, day_opt)

dts.day = get_day_of_month(&dts, day_opt)
out[i] = dtstruct_to_dt64(&dts)
# offset semantics - if on the anchor point and going backwards
# shift to next
if n <= 0 and (months_since != 0 or
(months_since == 0 and dts.day > compare_day)):
# make sure to roll forward, so negate
n += 1
elif n > 0 and (months_since == 0 and dts.day < compare_day):
# pretend to roll back if on same month but
# before compare_day
n -= 1

else:
raise ValueError("day must be None, 'start', 'end', "
"'business_start', or 'business_end'")
dts.year = year_add_months(dts, modby * n - months_since)
dts.month = month_add_months(dts, modby * n - months_since)
dts.day = get_day_of_month(&dts, day_opt)

return np.asarray(out)
out[i] = dtstruct_to_dt64(&dts)


cdef ndarray[int64_t] shift_bdays(const int64_t[:] i8other, int periods):
Expand Down Expand Up @@ -4035,8 +3908,7 @@ cdef ndarray[int64_t] shift_bdays(const int64_t[:] i8other, int periods):
return result.base


def shift_month(stamp: datetime, months: int,
day_opt: object=None) -> datetime:
def shift_month(stamp: datetime, months: int, day_opt: object=None) -> datetime:
"""
Given a datetime (or Timestamp) `stamp`, an integer `months` and an
option `day_opt`, return a new datetimelike that many months later,
Expand Down Expand Up @@ -4078,14 +3950,14 @@ def shift_month(stamp: datetime, months: int,
if day_opt is None:
days_in_month = get_days_in_month(year, month)
day = min(stamp.day, days_in_month)
elif day_opt == 'start':
elif day_opt == "start":
day = 1
elif day_opt == 'end':
elif day_opt == "end":
day = get_days_in_month(year, month)
elif day_opt == 'business_start':
elif day_opt == "business_start":
# first business day of month
day = get_firstbday(year, month)
elif day_opt == 'business_end':
elif day_opt == "business_end":
# last business day of month
day = get_lastbday(year, month)
elif is_integer_object(day_opt):
Expand Down Expand Up @@ -4126,15 +3998,15 @@ cdef inline int get_day_of_month(npy_datetimestruct* dts, day_opt) nogil except?
cdef:
int days_in_month

if day_opt == 'start':
if day_opt == "start":
Copy link
Contributor

Choose a reason for hiding this comment

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

maybe want to make these an enum

Copy link
Member Author

Choose a reason for hiding this comment

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

yah i looked at this, there are a couple of usages that dont follow this pattern yet

return 1
elif day_opt == 'end':
elif day_opt == "end":
days_in_month = get_days_in_month(dts.year, dts.month)
return days_in_month
elif day_opt == 'business_start':
elif day_opt == "business_start":
# first business day of month
return get_firstbday(dts.year, dts.month)
elif day_opt == 'business_end':
elif day_opt == "business_end":
# last business day of month
return get_lastbday(dts.year, dts.month)
elif day_opt is not None:
Expand Down