simplify algebra in Year offset apply methods (#18280)

Author: jbrockmendel

pandas/_libs/tslibs/offsets.pyx

@@ -428,3 +428,105 @@ cpdef datetime shift_month(datetime stamp, int months, object day_opt=None):
     else:
         raise ValueError(day_opt)
     return stamp.replace(year=year, month=month, day=day)
+
+
+cdef int get_day_of_month(datetime other, day_opt) except? -1:
+    """
+    Find the day in `other`'s month that satisfies a DateOffset's onOffset
+    policy, as described by the `day_opt` argument.
+
+    Parameters
+    ----------
+    other : datetime or Timestamp
+    day_opt : 'start', 'end'
+        'start': returns 1
+        'end': returns last day  of the month
+
+    Returns
+    -------
+    day_of_month : int
+
+    Examples
+    -------
+    >>> other = datetime(2017, 11, 14)
+    >>> get_day_of_month(other, 'start')
+    1
+    >>> get_day_of_month(other, 'end')
+    30
+
+    """
+    if day_opt == 'start':
+        return 1
+    elif day_opt == 'end':
+        return monthrange(other.year, other.month)[1]
+    else:
+        raise ValueError(day_opt)
+
+
+cpdef int roll_yearday(other, n, month, day_opt='start') 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
+    day_opt : 'start', 'end'
+        'start': returns 1
+        'end': returns last day  of the month
+
+    Returns
+    -------
+    n : int number of periods to increment
+
+    Notes
+    -----
+    * Mirrors `roll_check` in tslib.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
+
+    """
+    # 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(other,
+                                                                day_opt)):
+            n -= 1
+    elif n <= 0:
+        if other.month > month or (other.month == month and
+                                   other.day > get_day_of_month(other,
+                                                                day_opt)):
+            n += 1
+    return n

pandas/tests/tseries/offsets/test_liboffsets.py

@@ -0,0 +1,63 @@
+# -*- coding: utf-8 -*-
+"""
+Tests for helper functions in the cython tslibs.offsets
+"""
+from datetime import datetime
+
+import pytest
+
+import pandas as pd
+
+import pandas._libs.tslibs.offsets as liboffsets
+
+
+def test_shift_month():
+    dt = datetime(2017, 11, 15)
+
+    assert liboffsets.shift_month(dt, 0, day_opt=None) == dt
+    assert liboffsets.shift_month(dt, 0, day_opt=15) == dt
+
+    assert liboffsets.shift_month(dt, 1,
+                                  day_opt='start') == datetime(2017, 12, 1)
+
+    assert liboffsets.shift_month(dt, -145,
+                                  day_opt='end') == datetime(2005, 10, 31)
+
+    with pytest.raises(ValueError):
+        liboffsets.shift_month(dt, 3, day_opt='this should raise')
+
+
+def test_get_day_of_month():
+    # get_day_of_month is not directly exposed; we test it via roll_yearday
+    dt = datetime(2017, 11, 15)
+
+    with pytest.raises(ValueError):
+        # To hit the raising case we need month == dt.month and n > 0
+        liboffsets.roll_yearday(dt, n=3, month=11, day_opt='foo')
+
+
+def test_roll_yearday():
+    # Copied from doctest examples
+    month = 3
+    day_opt = 'start'              # `other` will be compared to March 1
+    other = datetime(2017, 2, 10)  # before March 1
+    assert liboffsets.roll_yearday(other, 2, month, day_opt) == 1
+    assert liboffsets.roll_yearday(other, -7, month, day_opt) == -7
+    assert liboffsets.roll_yearday(other, 0, month, day_opt) == 0
+
+    other = pd.Timestamp('2014-03-15', tz='US/Eastern')  # after March 1
+    assert liboffsets.roll_yearday(other, 2, month, day_opt) == 2
+    assert liboffsets.roll_yearday(other, -7, month, day_opt) == -6
+    assert liboffsets.roll_yearday(other, 0, month, day_opt) == 1
+
+    month = 6
+    day_opt = 'end'                # `other` will be compared to June 30
+    other = datetime(1999, 6, 29)  # before June 30
+    assert liboffsets.roll_yearday(other, 5, month, day_opt) == 4
+    assert liboffsets.roll_yearday(other, -7, month, day_opt) == -7
+    assert liboffsets.roll_yearday(other, 0, month, day_opt) == 0
+
+    other = pd.Timestamp(2072, 8, 24, 6, 17, 18)  # after June 30
+    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

pandas/tseries/offsets.py

@@ -22,6 +22,7 @@
     _int_to_weekday, _weekday_to_int,
     _determine_offset,
     apply_index_wraps,
+    roll_yearday,
     shift_month,
     BeginMixin, EndMixin,
     BaseOffset)
@@ -1905,49 +1906,12 @@ class YearEnd(EndMixin, YearOffset):
 
     @apply_wraps
     def apply(self, other):
-        def _increment(date):
-            if date.month == self.month:
-                _, days_in_month = tslib.monthrange(date.year, self.month)
-                if date.day != days_in_month:
-                    year = date.year
-                else:
-                    year = date.year + 1
-            elif date.month < self.month:
-                year = date.year
-            else:
-                year = date.year + 1
-            _, days_in_month = tslib.monthrange(year, self.month)
-            return datetime(year, self.month, days_in_month,
-                            date.hour, date.minute, date.second,
-                            date.microsecond)
-
-        def _decrement(date):
-            year = date.year if date.month > self.month else date.year - 1
-            _, days_in_month = tslib.monthrange(year, self.month)
-            return datetime(year, self.month, days_in_month,
-                            date.hour, date.minute, date.second,
-                            date.microsecond)
-
-        def _rollf(date):
-            if date.month != self.month or\
-               date.day < tslib.monthrange(date.year, date.month)[1]:
-                date = _increment(date)
-            return date
-
-        n = self.n
-        result = other
-        if n > 0:
-            while n > 0:
-                result = _increment(result)
-                n -= 1
-        elif n < 0:
-            while n < 0:
-                result = _decrement(result)
-                n += 1
-        else:
-            # n == 0, roll forward
-            result = _rollf(result)
-        return result
+        n = roll_yearday(other, self.n, self.month, 'end')
+        year = other.year + n
+        days_in_month = tslib.monthrange(year, self.month)[1]
+        return datetime(year, self.month, days_in_month,
+                        other.hour, other.minute, other.second,
+                        other.microsecond)
 
     @apply_index_wraps
     def apply_index(self, i):
@@ -1968,36 +1932,9 @@ class YearBegin(BeginMixin, YearOffset):
 
     @apply_wraps
     def apply(self, other):
-        def _increment(date, n):
-            year = date.year + n - 1
-            if date.month >= self.month:
-                year += 1
-            return datetime(year, self.month, 1, date.hour, date.minute,
-                            date.second, date.microsecond)
-
-        def _decrement(date, n):
-            year = date.year + n + 1
-            if date.month < self.month or (date.month == self.month and
-                                           date.day == 1):
-                year -= 1
-            return datetime(year, self.month, 1, date.hour, date.minute,
-                            date.second, date.microsecond)
-
-        def _rollf(date):
-            if (date.month != self.month) or date.day > 1:
-                date = _increment(date, 1)
-            return date
-
-        n = self.n
-        result = other
-        if n > 0:
-            result = _increment(result, n)
-        elif n < 0:
-            result = _decrement(result, n)
-        else:
-            # n == 0, roll forward
-            result = _rollf(result)
-        return result
+        n = roll_yearday(other, self.n, self.month, 'start')
+        year = other.year + n
+        return other.replace(year=year, month=self.month, day=1)
 
     @apply_index_wraps
     def apply_index(self, i):