simplify algebra in Year offset apply methods

Author: jbrockmendel

pandas/_libs/tslibs/offsets.pyx

@@ -428,3 +428,36 @@ 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)
+
+
+cpdef int _get_day_of_month(datetime other, day_opt):
+    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'):
+    """
+    Possibly increment or decrement the number of periods to shift
+    based on rollforward/rollbackward conventions.
+
+    Mirrors `roll_check` in tslib.shift_months
+    """
+    # 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/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)
@@ -1943,49 +1944,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):
@@ -2006,36 +1970,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):