Implement roll_monthday, simplify SemiMonthOffset

pandas/_libs/tslibs/offsets.pyx

@@ -523,11 +523,9 @@ def shift_quarters(int64_t[:] dtindex, int quarters,
                 n = quarters
 
                 months_since = (dts.month - q1start_month) % modby
-                compare_month = dts.month - months_since
-                compare_month = compare_month or 12
                 # compare_day is only relevant for comparison in the case
                 # where months_since == 0.
-                compare_day = get_firstbday(dts.year, compare_month)
+                compare_day = get_firstbday(dts.year, dts.month)
 
                 if n <= 0 and (months_since != 0 or
                                (months_since == 0 and dts.day > compare_day)):
@@ -556,11 +554,9 @@ def shift_quarters(int64_t[:] dtindex, int quarters,
                 n = quarters
 
                 months_since = (dts.month - q1start_month) % modby
-                compare_month = dts.month - months_since
-                compare_month = compare_month or 12
                 # compare_day is only relevant for comparison in the case
                 # where months_since == 0.
-                compare_day = get_lastbday(dts.year, compare_month)
+                compare_day = get_lastbday(dts.year, dts.month)
 
                 if n <= 0 and (months_since != 0 or
                                (months_since == 0 and dts.day > compare_day)):
@@ -827,7 +823,55 @@ cpdef int get_day_of_month(datetime other, day_opt) except? -1:
         raise ValueError(day_opt)
 
 
-cpdef int roll_yearday(other, n, month, day_opt='start') except? -1:
+cpdef int roll_convention(int other, int n, int compare):
+    """
+    Possibly increment or decrement the number of periods to shift
+    based on rollforward/rollbackward conventions.
+
+    Parameters
+    ----------
+    other : int, generally the day component of a datetime
+    n : number of periods to increment, before adjusting for rolling
+    compare : int, generally the day component of a datetime, in the same
+              month as the datetime form which `other` was taken.
+
+    Returns
+    -------
+    n : int number of periods to increment
+    """
+    if n > 0 and other < compare:
+        n -= 1
+    elif n <= 0 and other > compare:
+        # as if rolled forward already
+        n += 1
+    return n
+
+
+cpdef int roll_monthday(datetime other, int n, datetime compare):
+    """
+    Possibly increment or decrement the number of periods to shift
+    based on rollforward/rollbackward conventions.
+
+    Parameters
+    ----------
+    other : datetime
+    n : number of periods to increment, before adjusting for rolling
+    compare : datetime
+
+    Returns
+    -------
+    n : int number of periods to increment
+    """
+    if n > 0 and other < compare:
+        n -= 1
+    elif n <= 0 and other > compare:
+        # as if rolled forward already
+        n += 1
+    return n
+
+
+cpdef int roll_qtrday(datetime other, int n, int month, object day_opt,
+                      int modby=3) except? -1:
     """
     Possibly increment or decrement the number of periods to shift
     based on rollforward/rollbackward conventions.
@@ -836,6 +880,48 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1:
     ----------
     other : datetime or Timestamp
     n : number of periods to increment, before adjusting for rolling
+    month : int reference month giving the first month of the year
+    day_opt : 'start', 'end', 'business_start', 'business_end'
+        The convention to use in finding the day in a given month against
+        which to compare for rollforward/rollbackward decisions.
+    modby : int 3 for quarters, 12 for years
+
+    Returns
+    -------
+    n : int number of periods to increment
+    """
+    # 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 n > 0:
+        if months_since < 0 or (months_since == 0 and
+                                other.day < get_day_of_month(other,
+                                                             day_opt)):
+            # pretend to roll back if on same month but
+            # before compare_day
+            n -= 1
+    else:
+        if months_since > 0 or (months_since == 0 and
+                                other.day > get_day_of_month(other,
+                                                             day_opt)):
+            # make sure to roll forward, so negate
+            n += 1
+    return n
+
+
+cpdef int roll_yearday(datetime other, int n, int month,
+                       object day_opt) except? -1:
+    """
+    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'
         'start': returns 1
         'end': returns last day  of the month
@@ -846,7 +932,7 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1:
 
     Notes
     -----
-    * Mirrors `roll_check` in tslib.shift_months
+    * Mirrors `roll_check` in shift_months
 
     Examples
     -------
@@ -888,7 +974,7 @@ cpdef int roll_yearday(other, n, month, day_opt='start') except? -1:
                                    other.day < get_day_of_month(other,
                                                                 day_opt)):
             n -= 1
-    elif n <= 0:
+    else:
         if other.month > month or (other.month == month and
                                    other.day > get_day_of_month(other,
                                                                 day_opt)):

pandas/tests/tseries/offsets/test_liboffsets.py

@@ -9,6 +9,7 @@
 from pandas import Timestamp
 
 import pandas._libs.tslibs.offsets as liboffsets
+from pandas._libs.tslibs.offsets import roll_qtrday
 
 
 def test_get_lastbday():
@@ -95,3 +96,93 @@ def test_roll_yearday():
     assert liboffsets.roll_yearday(other, 5, month, day_opt) == 5
     assert liboffsets.roll_yearday(other, -7, month, day_opt) == -6
     assert liboffsets.roll_yearday(other, 0, month, day_opt) == 1
+
+
+def test_roll_qtrday():
+    other = Timestamp(2072, 10, 1, 6, 17, 18)  # Saturday
+    for day_opt in ['start', 'end', 'business_start', 'business_end']:
+        # as long as (other.month % 3) != (month % 3), day_opt is irrelevant
+        # the `day_opt` doesn't matter.
+        month = 5  # (other.month % 3) < (month % 3)
+        assert roll_qtrday(other, 4, month, day_opt, modby=3) == 3
+        assert roll_qtrday(other, -3, month, day_opt, modby=3) == -3
+
+        month = 3  # (other.month % 3) > (month % 3)
+        assert roll_qtrday(other, 4, month, day_opt, modby=3) == 4
+        assert roll_qtrday(other, -3, month, day_opt, modby=3) == -2
+
+    month = 2
+    other = datetime(1999, 5, 31)  # Monday
+    # has (other.month % 3) == (month % 3)
+
+    n = 2
+    assert roll_qtrday(other, n, month, 'start', modby=3) == n
+    assert roll_qtrday(other, n, month, 'end', modby=3) == n
+    assert roll_qtrday(other, n, month, 'business_start', modby=3) == n
+    assert roll_qtrday(other, n, month, 'business_end', modby=3) == n
+
+    n = -1
+    assert roll_qtrday(other, n, month, 'start', modby=3) == n + 1
+    assert roll_qtrday(other, n, month, 'end', modby=3) == n
+    assert roll_qtrday(other, n, month, 'business_start', modby=3) == n + 1
+    assert roll_qtrday(other, n, month, 'business_end', modby=3) == n
+
+    other = Timestamp(2072, 10, 1, 6, 17, 18)  # Saturday
+    month = 4  # (other.month % 3) == (month % 3)
+    n = 2
+    assert roll_qtrday(other, n, month, 'start', modby=3) == n
+    assert roll_qtrday(other, n, month, 'end', modby=3) == n - 1
+    assert roll_qtrday(other, n, month, 'business_start', modby=3) == n - 1
+    assert roll_qtrday(other, n, month, 'business_end', modby=3) == n - 1
+
+    n = -1
+    assert roll_qtrday(other, n, month, 'start', modby=3) == n
+    assert roll_qtrday(other, n, month, 'end', modby=3) == n
+    assert roll_qtrday(other, n, month, 'business_start', modby=3) == n
+    assert roll_qtrday(other, n, month, 'business_end', modby=3) == n
+
+    other = Timestamp(2072, 10, 3, 6, 17, 18)  # First businessday
+    month = 4  # (other.month % 3) == (month % 3)
+    n = 2
+    assert roll_qtrday(other, n, month, 'start', modby=3) == n
+    assert roll_qtrday(other, n, month, 'end', modby=3) == n - 1
+    assert roll_qtrday(other, n, month, 'business_start', modby=3) == n
+    assert roll_qtrday(other, n, month, 'business_end', modby=3) == n - 1
+
+    n = -1
+    assert roll_qtrday(other, n, month, 'start', modby=3) == n + 1
+    assert roll_qtrday(other, n, month, 'end', modby=3) == n
+    assert roll_qtrday(other, n, month, 'business_start', modby=3) == n
+    assert roll_qtrday(other, n, month, 'business_end', modby=3) == n
+
+
+def test_roll_monthday():
+    other = Timestamp('2017-12-29', tz='US/Pacific')
+    before = Timestamp('2017-12-01', tz='US/Pacific')
+    after = Timestamp('2017-12-31', tz='US/Pacific')
+
+    n = 42
+    assert liboffsets.roll_monthday(other, n, other) == n
+    assert liboffsets.roll_monthday(other, n, before) == n
+    assert liboffsets.roll_monthday(other, n, after) == n - 1
+
+    n = -4
+    assert liboffsets.roll_monthday(other, n, other) == n
+    assert liboffsets.roll_monthday(other, n, before) == n + 1
+    assert liboffsets.roll_monthday(other, n, after) == n
+
+
+def test_roll_convention():
+    other = 29
+    before = 1
+    after = 31
+
+    n = 42
+    assert liboffsets.roll_convention(other, n, other) == n
+    assert liboffsets.roll_convention(other, n, before) == n
+    assert liboffsets.roll_convention(other, n, after) == n - 1
+
+    n = -4
+    assert liboffsets.roll_convention(other, n, other) == n
+    assert liboffsets.roll_convention(other, n, before) == n + 1
+    assert liboffsets.roll_convention(other, n, after) == n