From df989d88a46794657f10b50c3aabe136f23f21f1 Mon Sep 17 00:00:00 2001 From: Brock Date: Sat, 13 Jun 2020 18:30:22 -0700 Subject: [PATCH 1/2] REF: implement shift_bday --- pandas/_libs/tslibs/offsets.pyx | 66 +++++++++++++++++++++++---------- 1 file changed, 47 insertions(+), 19 deletions(-) diff --git a/pandas/_libs/tslibs/offsets.pyx b/pandas/_libs/tslibs/offsets.pyx index 4069d192d9e88..558ce879ecf9f 100644 --- a/pandas/_libs/tslibs/offsets.pyx +++ b/pandas/_libs/tslibs/offsets.pyx @@ -18,7 +18,7 @@ from dateutil.easter import easter import numpy as np cimport numpy as cnp -from numpy cimport int64_t +from numpy cimport int64_t, ndarray cnp.import_array() # TODO: formalize having _libs.properties "above" tslibs in the dependency structure @@ -1373,24 +1373,7 @@ cdef class BusinessDay(BusinessMixin): @apply_index_wraps def apply_index(self, dtindex): i8other = dtindex.asi8 - time = (i8other % DAY_NANOS).view("timedelta64[ns]") - - # to_period rolls forward to next BDay; track and - # reduce n where it does when rolling forward - asper = dtindex.to_period("B") - - if self.n > 0: - shifted = (dtindex.to_perioddelta("B") - time).asi8 != 0 - - roll = np.where(shifted, self.n - 1, self.n) - shifted = asper._addsub_int_array(roll, operator.add) - else: - # Integer addition is deprecated, so we use _time_shift directly - roll = self.n - shifted = asper._time_shift(roll) - - result = shifted.to_timestamp() + time - return result + return shift_bdays(i8other, self.n) def is_on_offset(self, dt) -> bool: if self.normalize and not _is_normalized(dt): @@ -3991,6 +3974,51 @@ def shift_months(const int64_t[:] dtindex, int months, object day=None): return np.asarray(out) +cdef ndarray[int64_t] shift_bdays(const int64_t[:] i8other, int periods): + cdef: + Py_ssize_t i, n = len(i8other) + int64_t[:] result = np.empty(n, dtype="i8") + int64_t val, res + int wday, nadj, days + npy_datetimestruct dts + + for i in range(n): + val = i8other[i] + if val == NPY_NAT: + result[i] = NPY_NAT + else: + # The rest of this is effectively a copy of BusinessDay.apply + nadj = periods + weeks = nadj // 5 + dt64_to_dtstruct(val, &dts) + wday = dayofweek(dts.year, dts.month, dts.day) + + if nadj <= 0 and wday > 4: + # roll forward + nadj += 1 + + nadj -= 5 * weeks + + # nadj is always >= 0 at this point + if nadj == 0 and wday > 4: + # roll back + days = 4 - wday + elif wday > 4: + # roll forward + days = (7 - wday) + (nadj - 1) + elif wday + nadj <= 4: + # shift by n days without leaving the current week + days = nadj + else: + # shift by nadj days plus 2 to get past the weekend + days = nadj + 2 + + res = val + (7 * weeks + days) * DAY_NANOS + result[i] = res + + return result.base + + def shift_month(stamp: datetime, months: int, day_opt: object=None) -> datetime: """ From 32859a5175f05af99996700d56e97e9726c2edc6 Mon Sep 17 00:00:00 2001 From: Brock Date: Sun, 14 Jun 2020 12:44:18 -0700 Subject: [PATCH 2/2] docstring --- pandas/_libs/tslibs/offsets.pyx | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/pandas/_libs/tslibs/offsets.pyx b/pandas/_libs/tslibs/offsets.pyx index d4fa97ed11cdb..15de4332d9992 100644 --- a/pandas/_libs/tslibs/offsets.pyx +++ b/pandas/_libs/tslibs/offsets.pyx @@ -3974,6 +3974,18 @@ def shift_months(const int64_t[:] dtindex, int months, object day_opt=None): cdef ndarray[int64_t] shift_bdays(const int64_t[:] i8other, int periods): + """ + Implementation of BusinessDay.apply_offset. + + Parameters + ---------- + i8other : const int64_t[:] + periods : int + + Returns + ------- + ndarray[int64_t] + """ cdef: Py_ssize_t i, n = len(i8other) int64_t[:] result = np.empty(n, dtype="i8")