move and de-privatize _localize_pydatetime, move shift_day to liboffsets

pandas/_libs/tslib.pyx

@@ -54,6 +54,7 @@ from tslibs.nattype cimport checknull_with_nat, NPY_NAT
 from tslibs.timestamps cimport (create_timestamp_from_ts,
                                 _NS_UPPER_BOUND, _NS_LOWER_BOUND)
 from tslibs.timestamps import Timestamp
+from tslibs.offsets import localize_pydatetime  # noqa:F841
 
 cdef bint PY2 = str == bytes
 
@@ -234,23 +235,6 @@ def _test_parse_iso8601(object ts):
         return Timestamp(obj.value)
 
 
-cpdef inline object _localize_pydatetime(object dt, object tz):
-    """
-    Take a datetime/Timestamp in UTC and localizes to timezone tz.
-    """
-    if tz is None:
-        return dt
-    elif isinstance(dt, Timestamp):
-        return dt.tz_localize(tz)
-    elif tz == 'UTC' or tz is UTC:
-        return UTC.localize(dt)
-    try:
-        # datetime.replace with pytz may be incorrect result
-        return tz.localize(dt)
-    except AttributeError:
-        return dt.replace(tzinfo=tz)
-
-
 def format_array_from_datetime(ndarray[int64_t] values, object tz=None,
                                object format=None, object na_rep=None):
     """

pandas/_libs/tslibs/offsets.pyx

@@ -5,9 +5,13 @@ cimport cython
 from cython cimport Py_ssize_t
 
 import time
-from cpython.datetime cimport datetime, timedelta, time as dt_time
+from cpython.datetime cimport (PyDateTime_IMPORT, PyDateTime_CheckExact,
+                               datetime, timedelta,
+                               time as dt_time)
+PyDateTime_IMPORT
 
 from dateutil.relativedelta import relativedelta
+from pytz import UTC
 
 import numpy as np
 cimport numpy as cnp
@@ -494,6 +498,58 @@ class BaseOffset(_BaseOffset):
 # ----------------------------------------------------------------------
 # RelativeDelta Arithmetic
 
+cpdef inline datetime localize_pydatetime(datetime dt, object tz):
+    """
+    Take a datetime/Timestamp in UTC and localizes to timezone tz.
+
+    Parameters
+    ----------
+    dt : datetime or Timestamp
+    tz : tzinfo, "UTC", or None
+
+    Returns
+    -------
+    localized : datetime or Timestamp
+    """
+    if tz is None:
+        return dt
+    elif not PyDateTime_CheckExact(dt):
+        # i.e. is a Timestamp
+        return dt.tz_localize(tz)
+    elif tz == 'UTC' or tz is UTC:
+        return UTC.localize(dt)
+    try:
+        # datetime.replace with pytz may be incorrect result
+        return tz.localize(dt)
+    except AttributeError:
+        return dt.replace(tzinfo=tz)
+
+
+cpdef datetime shift_day(datetime other, int days):
+    """
+    Increment the datetime `other` by the given number of days, retaining
+    the time-portion of the datetime.  For tz-naive datetimes this is
+    equivalent to adding a timedelta.  For tz-aware datetimes it is similar to
+    dateutil's relativedelta.__add__, but handles pytz tzinfo objects.
+
+    Parameters
+    ----------
+    other : datetime or Timestamp
+    days : int
+
+    Returns
+    -------
+    shifted: datetime or Timestamp
+    """
+    if other.tzinfo is None:
+        return other + timedelta(days=days)
+
+    tz = other.tzinfo
+    naive = other.replace(tzinfo=None)
+    shifted = naive + timedelta(days=days)
+    return localize_pydatetime(shifted, tz)
+
+
 cdef inline int year_add_months(pandas_datetimestruct dts, int months) nogil:
     """new year number after shifting pandas_datetimestruct number of months"""
     return dts.year + (dts.month + months - 1) / 12

pandas/tests/indexes/datetimes/test_arithmetic.py

@@ -56,10 +56,7 @@ def test_dti_cmp_datetimelike(self, other, tz):
             if isinstance(other, np.datetime64):
                 # no tzaware version available
                 return
-            elif isinstance(other, Timestamp):
-                other = other.tz_localize(dti.tzinfo)
-            else:
-                other = tslib._localize_pydatetime(other, dti.tzinfo)
+            other = tslib.localize_pydatetime(other, dti.tzinfo)
 
         result = dti == other
         expected = np.array([True, False])

pandas/tests/indexes/datetimes/test_timezones.py

@@ -911,12 +911,12 @@ def test_with_tz(self, tz):
         central = dr.tz_convert(tz)
         assert central.tz is tz
         naive = central[0].to_pydatetime().replace(tzinfo=None)
-        comp = tslib._localize_pydatetime(naive, tz).tzinfo
+        comp = tslib.localize_pydatetime(naive, tz).tzinfo
         assert central[0].tz is comp
 
         # compare vs a localized tz
         naive = dr[0].to_pydatetime().replace(tzinfo=None)
-        comp = tslib._localize_pydatetime(naive, tz).tzinfo
+        comp = tslib.localize_pydatetime(naive, tz).tzinfo
         assert central[0].tz is comp
 
         # datetimes with tzinfo set
@@ -946,7 +946,7 @@ def test_dti_convert_tz_aware_datetime_datetime(self, tz):
         dates = [datetime(2000, 1, 1), datetime(2000, 1, 2),
                  datetime(2000, 1, 3)]
 
-        dates_aware = [tslib._localize_pydatetime(x, tz) for x in dates]
+        dates_aware = [tslib.localize_pydatetime(x, tz) for x in dates]
         result = DatetimeIndex(dates_aware)
         assert timezones.tz_compare(result.tz, tz)
 

pandas/tests/scalar/timestamp/test_unary_ops.py

@@ -242,7 +242,7 @@ def test_replace_across_dst(self, tz, normalize):
         # GH#18319 check that 1) timezone is correctly normalized and
         # 2) that hour is not incorrectly changed by this normalization
         ts_naive = Timestamp('2017-12-03 16:03:30')
-        ts_aware = tslib._localize_pydatetime(ts_naive, tz)
+        ts_aware = tslib.localize_pydatetime(ts_naive, tz)
 
         # Preliminary sanity-check
         assert ts_aware == normalize(ts_aware)

pandas/tests/series/test_timezones.py

@@ -298,7 +298,7 @@ def test_getitem_pydatetime_tz(self, tzstr):
         time_pandas = Timestamp('2012-12-24 17:00', tz=tzstr)
 
         dt = datetime(2012, 12, 24, 17, 0)
-        time_datetime = tslib._localize_pydatetime(dt, tz)
+        time_datetime = tslib.localize_pydatetime(dt, tz)
         assert ts[time_pandas] == ts[time_datetime]
 
     def test_series_truncate_datetimeindex_tz(self):

pandas/tests/tseries/offsets/test_offsets.py

@@ -319,7 +319,7 @@ def _check_offsetfunc_works(self, offset, funcname, dt, expected,
         for tz in self.timezones:
             expected_localize = expected.tz_localize(tz)
             tz_obj = timezones.maybe_get_tz(tz)
-            dt_tz = tslib._localize_pydatetime(dt, tz_obj)
+            dt_tz = tslib.localize_pydatetime(dt, tz_obj)
 
             result = func(dt_tz)
             assert isinstance(result, Timestamp)

pandas/tests/tslibs/test_timezones.py

@@ -51,17 +51,17 @@ def test_infer_tz(eastern, localize):
     end = localize(eastern, end_naive)
 
     assert (timezones.infer_tzinfo(start, end) is
-            tslib._localize_pydatetime(start_naive, eastern).tzinfo)
+            tslib.localize_pydatetime(start_naive, eastern).tzinfo)
     assert (timezones.infer_tzinfo(start, None) is
-            tslib._localize_pydatetime(start_naive, eastern).tzinfo)
+            tslib.localize_pydatetime(start_naive, eastern).tzinfo)
     assert (timezones.infer_tzinfo(None, end) is
-            tslib._localize_pydatetime(end_naive, eastern).tzinfo)
+            tslib.localize_pydatetime(end_naive, eastern).tzinfo)
 
     start = utc.localize(start_naive)
     end = utc.localize(end_naive)
     assert timezones.infer_tzinfo(start, end) is utc
 
-    end = tslib._localize_pydatetime(end_naive, eastern)
+    end = tslib.localize_pydatetime(end_naive, eastern)
     with pytest.raises(Exception):
         timezones.infer_tzinfo(start, end)
     with pytest.raises(Exception):

pandas/tseries/offsets.py

@@ -76,7 +76,7 @@ def wrapper(self, other):
             result = func(self, other)
 
             if self._adjust_dst:
-                result = tslib._localize_pydatetime(result, tz)
+                result = liboffsets.localize_pydatetime(result, tz)
 
             result = Timestamp(result)
             if self.normalize:
@@ -94,7 +94,7 @@ def wrapper(self, other):
                     result = Timestamp(value + nano)
 
             if tz is not None and result.tzinfo is None:
-                result = tslib._localize_pydatetime(result, tz)
+                result = liboffsets.localize_pydatetime(result, tz)
 
         except OutOfBoundsDatetime:
             result = func(self, as_datetime(other))
@@ -104,37 +104,12 @@ def wrapper(self, other):
                 result = tslib.normalize_date(result)
 
             if tz is not None and result.tzinfo is None:
-                result = tslib._localize_pydatetime(result, tz)
+                result = liboffsets.localize_pydatetime(result, tz)
 
         return result
     return wrapper
 
 
-def shift_day(other, days):
-    """
-    Increment the datetime `other` by the given number of days, retaining
-    the time-portion of the datetime.  For tz-naive datetimes this is
-    equivalent to adding a timedelta.  For tz-aware datetimes it is similar to
-    dateutil's relativedelta.__add__, but handles pytz tzinfo objects.
-
-    Parameters
-    ----------
-    other : datetime or Timestamp
-    days : int
-
-    Returns
-    -------
-    shifted: datetime or Timestamp
-    """
-    if other.tzinfo is None:
-        return other + timedelta(days=days)
-
-    tz = other.tzinfo
-    naive = other.replace(tzinfo=None)
-    shifted = naive + timedelta(days=days)
-    return tslib._localize_pydatetime(shifted, tz)
-
-
 # ---------------------------------------------------------------------
 # DateOffset
 
@@ -221,7 +196,7 @@ def apply(self, other):
 
             if tzinfo is not None and self._use_relativedelta:
                 # bring tz back from UTC calculation
-                other = tslib._localize_pydatetime(other, tzinfo)
+                other = liboffsets.localize_pydatetime(other, tzinfo)
 
             return as_timestamp(other)
         else:
@@ -1355,7 +1330,7 @@ def apply(self, other):
 
         shifted = shift_month(other, months, 'start')
         to_day = self._get_offset_day(shifted)
-        return shift_day(shifted, to_day - shifted.day)
+        return liboffsets.shift_day(shifted, to_day - shifted.day)
 
     def onOffset(self, dt):
         if self.normalize and not _is_normalized(dt):
@@ -1781,9 +1756,9 @@ def apply(self, other):
         next_year = self.get_year_end(
             datetime(other.year + 1, self.startingMonth, 1))
 
-        prev_year = tslib._localize_pydatetime(prev_year, other.tzinfo)
-        cur_year = tslib._localize_pydatetime(cur_year, other.tzinfo)
-        next_year = tslib._localize_pydatetime(next_year, other.tzinfo)
+        prev_year = liboffsets.localize_pydatetime(prev_year, other.tzinfo)
+        cur_year = liboffsets.localize_pydatetime(cur_year, other.tzinfo)
+        next_year = liboffsets.localize_pydatetime(next_year, other.tzinfo)
 
         # Note: next_year.year == other.year + 1, so we will always
         # have other < next_year
@@ -1984,7 +1959,7 @@ def _rollback_to_year(self, other):
             qtr_lens = self.get_weeks(norm)
 
             # check thet qtr_lens is consistent with self._offset addition
-            end = shift_day(start, days=7 * sum(qtr_lens))
+            end = liboffsets.shift_day(start, days=7 * sum(qtr_lens))
             assert self._offset.onOffset(end), (start, end, qtr_lens)
 
             tdelta = norm - start
@@ -2024,7 +1999,7 @@ def apply(self, other):
         # Note: we always have 0 <= n < 4
         weeks = sum(qtr_lens[:n])
         if weeks:
-            res = shift_day(res, days=weeks * 7)
+            res = liboffsets.shift_day(res, days=weeks * 7)
 
         return res
 
@@ -2061,7 +2036,7 @@ def onOffset(self, dt):
 
         current = next_year_end
         for qtr_len in qtr_lens:
-            current = shift_day(current, days=qtr_len * 7)
+            current = liboffsets.shift_day(current, days=qtr_len * 7)
             if dt == current:
                 return True
         return False
@@ -2096,8 +2071,8 @@ def apply(self, other):
         current_easter = easter(other.year)
         current_easter = datetime(current_easter.year,
                                   current_easter.month, current_easter.day)
-        current_easter = tslib._localize_pydatetime(current_easter,
-                                                    other.tzinfo)
+        current_easter = liboffsets.localize_pydatetime(current_easter,
+                                                        other.tzinfo)
 
         n = self.n
         if n >= 0 and other < current_easter: