REF: re-use tz_convert_from_utc_single in _localize_tso (pandas-dev#46525)

Author: jbrockmendel

pandas/_libs/tslibs/__init__.py

@@ -55,7 +55,9 @@
 )
 from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._libs.tslibs.timezones import tz_compare
-from pandas._libs.tslibs.tzconversion import tz_convert_from_utc_single
+from pandas._libs.tslibs.tzconversion import (
+    py_tz_convert_from_utc_single as tz_convert_from_utc_single,
+)
 from pandas._libs.tslibs.vectorized import (
     dt64arr_to_periodarr,
     get_resolution,

pandas/_libs/tslibs/conversion.pyx

@@ -72,8 +72,9 @@ from pandas._libs.tslibs.nattype cimport (
 )
 from pandas._libs.tslibs.tzconversion cimport (
     bisect_right_i8,
-    infer_datetuil_fold,
+    infer_dateutil_fold,
     localize_tzinfo_api,
+    tz_convert_from_utc_single,
     tz_localize_to_utc_single,
 )
 
@@ -531,7 +532,7 @@ cdef _TSObject _create_tsobject_tz_using_offset(npy_datetimestruct dts,
         if typ == 'dateutil':
             tdata = <int64_t*>cnp.PyArray_DATA(trans)
             pos = bisect_right_i8(tdata, obj.value, trans.shape[0]) - 1
-            obj.fold = infer_datetuil_fold(obj.value, trans, deltas, pos)
+            obj.fold = infer_dateutil_fold(obj.value, trans, deltas, pos)
 
     # Keep the converter same as PyDateTime's
     dt = datetime(obj.dts.year, obj.dts.month, obj.dts.day,
@@ -683,7 +684,7 @@ cdef inline void _localize_tso(_TSObject obj, tzinfo tz):
         int64_t[::1] deltas
         int64_t local_val
         int64_t* tdata
-        Py_ssize_t pos, ntrans
+        Py_ssize_t pos, ntrans, outpos = -1
         str typ
 
     assert obj.tzinfo is None
@@ -692,35 +693,12 @@ cdef inline void _localize_tso(_TSObject obj, tzinfo tz):
         pass
     elif obj.value == NPY_NAT:
         pass
-    elif is_tzlocal(tz):
-        local_val = obj.value + localize_tzinfo_api(obj.value, tz, &obj.fold)
-        dt64_to_dtstruct(local_val, &obj.dts)
     else:
-        # Adjust datetime64 timestamp, recompute datetimestruct
-        trans, deltas, typ = get_dst_info(tz)
-        ntrans = trans.shape[0]
-
-        if typ == "pytz":
-            # i.e. treat_tz_as_pytz(tz)
-            tdata = <int64_t*>cnp.PyArray_DATA(trans)
-            pos = bisect_right_i8(tdata, obj.value, ntrans) - 1
-            local_val = obj.value + deltas[pos]
-
-            # find right representation of dst etc in pytz timezone
-            tz = tz._tzinfos[tz._transition_info[pos]]
-        elif typ == "dateutil":
-            # i.e. treat_tz_as_dateutil(tz)
-            tdata = <int64_t*>cnp.PyArray_DATA(trans)
-            pos = bisect_right_i8(tdata, obj.value, ntrans) - 1
-            local_val = obj.value + deltas[pos]
+        local_val = tz_convert_from_utc_single(obj.value, tz, &obj.fold, &outpos)
 
-            # dateutil supports fold, so we infer fold from value
-            obj.fold = infer_datetuil_fold(obj.value, trans, deltas, pos)
-        else:
-            # All other cases have len(deltas) == 1. As of 2018-07-17
-            #  (and 2022-03-07), all test cases that get here have
-            #  is_fixed_offset(tz).
-            local_val = obj.value + deltas[0]
+        if outpos != -1:
+            # infer we went through a pytz path
+            tz = tz._tzinfos[tz._transition_info[outpos]]
 
         dt64_to_dtstruct(local_val, &obj.dts)
 

pandas/_libs/tslibs/tzconversion.pxd

@@ -8,14 +8,16 @@ from numpy cimport (
 cdef int64_t localize_tzinfo_api(
     int64_t utc_val, tzinfo tz, bint* fold=*
 ) except? -1
-cpdef int64_t tz_convert_from_utc_single(int64_t val, tzinfo tz)
+cdef int64_t tz_convert_from_utc_single(
+    int64_t utc_val, tzinfo tz, bint* fold=?, Py_ssize_t* outpos=?
+) except? -1
 cdef int64_t tz_localize_to_utc_single(
     int64_t val, tzinfo tz, object ambiguous=*, object nonexistent=*
 ) except? -1
 
 cdef Py_ssize_t bisect_right_i8(int64_t *data, int64_t val, Py_ssize_t n)
 
-cdef bint infer_datetuil_fold(
+cdef bint infer_dateutil_fold(
     int64_t value,
     const int64_t[::1] trans,
     const int64_t[::1] deltas,

pandas/_libs/tslibs/tzconversion.pyi

@@ -12,7 +12,9 @@ def tz_convert_from_utc(
     vals: npt.NDArray[np.int64],  # const int64_t[:]
     tz: tzinfo,
 ) -> npt.NDArray[np.int64]: ...
-def tz_convert_from_utc_single(val: np.int64, tz: tzinfo) -> np.int64: ...
+
+# py_tz_convert_from_utc_single exposed for testing
+def py_tz_convert_from_utc_single(val: np.int64, tz: tzinfo) -> np.int64: ...
 def tz_localize_to_utc(
     vals: npt.NDArray[np.int64],
     tz: tzinfo | None,

pandas/_libs/tslibs/tzconversion.pyx

@@ -444,7 +444,18 @@ cdef int64_t localize_tzinfo_api(
     return _tz_localize_using_tzinfo_api(utc_val, tz, to_utc=False, fold=fold)
 
 
-cpdef int64_t tz_convert_from_utc_single(int64_t utc_val, tzinfo tz):
+def py_tz_convert_from_utc_single(int64_t utc_val, tzinfo tz):
+    # The 'bint* fold=NULL' in tz_convert_from_utc_single means we cannot
+    #  make it cdef, so this is version exposed for testing from python.
+    return tz_convert_from_utc_single(utc_val, tz)
+
+
+cdef int64_t tz_convert_from_utc_single(
+    int64_t utc_val,
+    tzinfo tz,
+    bint* fold=NULL,
+    Py_ssize_t* outpos=NULL,
+) except? -1:
     """
     Convert the val (in i8) from UTC to tz
 
@@ -454,6 +465,8 @@ cpdef int64_t tz_convert_from_utc_single(int64_t utc_val, tzinfo tz):
     ----------
     utc_val : int64
     tz : tzinfo
+    fold : bint*, default NULL
+    outpos : Py_ssize_t*, default NULL
 
     Returns
     -------
@@ -473,15 +486,31 @@ cpdef int64_t tz_convert_from_utc_single(int64_t utc_val, tzinfo tz):
         return utc_val
     elif is_tzlocal(tz):
         return utc_val + _tz_localize_using_tzinfo_api(utc_val, tz, to_utc=False)
-    elif is_fixed_offset(tz):
-        _, deltas, _ = get_dst_info(tz)
-        delta = deltas[0]
-        return utc_val + delta
     else:
-        trans, deltas, _ = get_dst_info(tz)
+        trans, deltas, typ = get_dst_info(tz)
         tdata = <int64_t*>cnp.PyArray_DATA(trans)
-        pos = bisect_right_i8(tdata, utc_val, trans.shape[0]) - 1
-        return utc_val + deltas[pos]
+
+        if typ == "dateutil":
+            pos = bisect_right_i8(tdata, utc_val, trans.shape[0]) - 1
+
+            if fold is not NULL:
+                fold[0] = infer_dateutil_fold(utc_val, trans, deltas, pos)
+            return utc_val + deltas[pos]
+
+        elif typ == "pytz":
+            pos = bisect_right_i8(tdata, utc_val, trans.shape[0]) - 1
+
+            # We need to get 'pos' back to the caller so it can pick the
+            #  correct "standardized" tzinfo objecg.
+            if outpos is not NULL:
+                outpos[0] = pos
+            return utc_val + deltas[pos]
+
+        else:
+            # All other cases have len(deltas) == 1. As of 2018-07-17
+            #  (and 2022-03-07), all test cases that get here have
+            #  is_fixed_offset(tz).
+            return utc_val + deltas[0]
 
 
 def tz_convert_from_utc(const int64_t[:] vals, tzinfo tz):
@@ -635,7 +664,7 @@ cdef int64_t _tz_localize_using_tzinfo_api(
 
 
 # NB: relies on dateutil internals, subject to change.
-cdef bint infer_datetuil_fold(
+cdef bint infer_dateutil_fold(
     int64_t value,
     const int64_t[::1] trans,
     const int64_t[::1] deltas,

pandas/tests/tslibs/test_conversion.py

@@ -21,7 +21,7 @@
 
 def _compare_utc_to_local(tz_didx):
     def f(x):
-        return tzconversion.tz_convert_from_utc_single(x, tz_didx.tz)
+        return tzconversion.py_tz_convert_from_utc_single(x, tz_didx.tz)
 
     result = tzconversion.tz_convert_from_utc(tz_didx.asi8, tz_didx.tz)
     expected = np.vectorize(f)(tz_didx.asi8)