ENH: support zoneinfo tzinfos (#46425)

* ENH: support zoneinfo tzinfos

* add backports.zoneinfo to ci deps

* add backports.zoneinfo to pypy file

* py38 compat

* fix zoneinfo check

* fix check on py38

* mypy fixup

* mypy fixup

* fix tznaive acse

* fix fold

* whatnsew

* flesh out comment

Author: jbrockmendel

doc/source/whatsnew/v1.5.0.rst

@@ -481,7 +481,7 @@ Timedelta
 
 Time Zones
 ^^^^^^^^^^
--
+- Bug in :class:`Timestamp` constructor raising when passed a ``ZoneInfo`` tzinfo object (:issue:`46425`)
 -
 
 Numeric

pandas/_libs/tslibs/conversion.pyx

@@ -56,6 +56,7 @@ from pandas._libs.tslibs.timezones cimport (
     is_fixed_offset,
     is_tzlocal,
     is_utc,
+    is_zoneinfo,
     maybe_get_tz,
     tz_compare,
     utc_pytz as UTC,
@@ -532,7 +533,7 @@ cdef _TSObject _create_tsobject_tz_using_offset(npy_datetimestruct dts,
     # see PEP 495 https://www.python.org/dev/peps/pep-0495/#the-fold-attribute
     if is_utc(tz):
         pass
-    elif is_tzlocal(tz):
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
         localize_tzinfo_api(obj.value, tz, &obj.fold)
     else:
         trans, deltas, typ = get_dst_info(tz)

pandas/_libs/tslibs/timezones.pxd

@@ -9,6 +9,7 @@ cdef tzinfo utc_pytz
 
 cpdef bint is_utc(tzinfo tz)
 cdef bint is_tzlocal(tzinfo tz)
+cdef bint is_zoneinfo(tzinfo tz)
 
 cdef bint treat_tz_as_pytz(tzinfo tz)
 

pandas/_libs/tslibs/timezones.pyx

@@ -3,6 +3,14 @@ from datetime import (
     timezone,
 )
 
+try:
+    # py39+
+    import zoneinfo
+    from zoneinfo import ZoneInfo
+except ImportError:
+    zoneinfo = None
+    ZoneInfo = None
+
 from cpython.datetime cimport (
     datetime,
     timedelta,
@@ -42,18 +50,43 @@ cdef tzinfo utc_stdlib = timezone.utc
 cdef tzinfo utc_pytz = UTC
 cdef tzinfo utc_dateutil_str = dateutil_gettz("UTC")  # NB: *not* the same as tzutc()
 
+cdef tzinfo utc_zoneinfo = None
+
 
 # ----------------------------------------------------------------------
 
+cdef inline bint is_utc_zoneinfo(tzinfo tz):
+    # Workaround for cases with missing tzdata
+    #  https://github.com/pandas-dev/pandas/pull/46425#discussion_r830633025
+    if tz is None or zoneinfo is None:
+        return False
+
+    global utc_zoneinfo
+    if utc_zoneinfo is None:
+        try:
+            utc_zoneinfo = ZoneInfo("UTC")
+        except zoneinfo.ZoneInfoNotFoundError:
+            return False
+
+    return tz is utc_zoneinfo
+
+
 cpdef inline bint is_utc(tzinfo tz):
     return (
         tz is utc_pytz
         or tz is utc_stdlib
         or isinstance(tz, _dateutil_tzutc)
         or tz is utc_dateutil_str
+        or is_utc_zoneinfo(tz)
     )
 
 
+cdef inline bint is_zoneinfo(tzinfo tz):
+    if ZoneInfo is None:
+        return False
+    return isinstance(tz, ZoneInfo)
+
+
 cdef inline bint is_tzlocal(tzinfo tz):
     return isinstance(tz, _dateutil_tzlocal)
 
@@ -210,6 +243,8 @@ cdef inline bint is_fixed_offset(tzinfo tz):
             return 1
         else:
             return 0
+    elif is_zoneinfo(tz):
+        return 0
     # This also implicitly accepts datetime.timezone objects which are
     # considered fixed
     return 1
@@ -264,6 +299,8 @@ cdef object get_dst_info(tzinfo tz):
         # e.g. pytz.FixedOffset, matplotlib.dates._UTC,
         # psycopg2.tz.FixedOffsetTimezone
         num = int(get_utcoffset(tz, None).total_seconds()) * 1_000_000_000
+        # If we have e.g. ZoneInfo here, the get_utcoffset call will return None,
+        #  so the total_seconds() call will raise AttributeError.
         return (np.array([NPY_NAT + 1], dtype=np.int64),
                 np.array([num], dtype=np.int64),
                 "unknown")
@@ -291,13 +328,13 @@ cdef object get_dst_info(tzinfo tz):
                 # deltas
                 deltas = np.array([v.offset for v in (
                     tz._ttinfo_before,) + tz._trans_idx], dtype='i8')
-                deltas *= 1000000000
+                deltas *= 1_000_000_000
                 typ = 'dateutil'
 
             elif is_fixed_offset(tz):
                 trans = np.array([NPY_NAT + 1], dtype=np.int64)
                 deltas = np.array([tz._ttinfo_std.offset],
-                                  dtype='i8') * 1000000000
+                                  dtype='i8') * 1_000_000_000
                 typ = 'fixed'
             else:
                 # 2018-07-12 this is not reached in the tests, and this case

pandas/_libs/tslibs/tzconversion.pyx

@@ -42,6 +42,7 @@ from pandas._libs.tslibs.timezones cimport (
     is_fixed_offset,
     is_tzlocal,
     is_utc,
+    is_zoneinfo,
     utc_pytz,
 )
 
@@ -60,7 +61,7 @@ cdef int64_t tz_localize_to_utc_single(
     elif is_utc(tz) or tz is None:
         return val
 
-    elif is_tzlocal(tz):
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
         return val - _tz_localize_using_tzinfo_api(val, tz, to_utc=True)
 
     elif is_fixed_offset(tz):
@@ -135,7 +136,7 @@ timedelta-like}
 
     result = np.empty(n, dtype=np.int64)
 
-    if is_tzlocal(tz):
+    if is_tzlocal(tz) or is_zoneinfo(tz):
         for i in range(n):
             v = vals[i]
             if v == NPY_NAT:
@@ -484,8 +485,8 @@ cdef int64_t tz_convert_from_utc_single(
 
     if is_utc(tz):
         return utc_val
-    elif is_tzlocal(tz):
-        return utc_val + _tz_localize_using_tzinfo_api(utc_val, tz, to_utc=False)
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
+        return utc_val + _tz_localize_using_tzinfo_api(utc_val, tz, to_utc=False, fold=fold)
     else:
         trans, deltas, typ = get_dst_info(tz)
         tdata = <int64_t*>cnp.PyArray_DATA(trans)
@@ -569,7 +570,7 @@ cdef const int64_t[:] _tz_convert_from_utc(const int64_t[:] stamps, tzinfo tz):
 
     if is_utc(tz) or tz is None:
         use_utc = True
-    elif is_tzlocal(tz):
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
         use_tzlocal = True
     else:
         trans, deltas, typ = get_dst_info(tz)

pandas/_libs/tslibs/vectorized.pyx

@@ -40,6 +40,7 @@ from .timezones cimport (
     get_dst_info,
     is_tzlocal,
     is_utc,
+    is_zoneinfo,
 )
 from .tzconversion cimport (
     bisect_right_i8,
@@ -117,7 +118,7 @@ def ints_to_pydatetime(
 
     if is_utc(tz) or tz is None:
         use_utc = True
-    elif is_tzlocal(tz):
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
         use_tzlocal = True
     else:
         trans, deltas, typ = get_dst_info(tz)
@@ -204,7 +205,7 @@ def get_resolution(const int64_t[:] stamps, tzinfo tz=None) -> Resolution:
 
     if is_utc(tz) or tz is None:
         use_utc = True
-    elif is_tzlocal(tz):
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
         use_tzlocal = True
     else:
         trans, deltas, typ = get_dst_info(tz)
@@ -272,7 +273,7 @@ cpdef ndarray[int64_t] normalize_i8_timestamps(const int64_t[:] stamps, tzinfo t
 
     if is_utc(tz) or tz is None:
         use_utc = True
-    elif is_tzlocal(tz):
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
         use_tzlocal = True
     else:
         trans, deltas, typ = get_dst_info(tz)
@@ -334,7 +335,7 @@ def is_date_array_normalized(const int64_t[:] stamps, tzinfo tz=None) -> bool:
 
     if is_utc(tz) or tz is None:
         use_utc = True
-    elif is_tzlocal(tz):
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
         use_tzlocal = True
     else:
         trans, deltas, typ = get_dst_info(tz)
@@ -385,7 +386,7 @@ def dt64arr_to_periodarr(const int64_t[:] stamps, int freq, tzinfo tz):
 
     if is_utc(tz) or tz is None:
         use_utc = True
-    elif is_tzlocal(tz):
+    elif is_tzlocal(tz) or is_zoneinfo(tz):
         use_tzlocal = True
     else:
         trans, deltas, typ = get_dst_info(tz)

pandas/conftest.py

@@ -75,6 +75,11 @@
     del pa
     has_pyarrow = True
 
+zoneinfo = None
+if pd.compat.PY39:
+    # Import "zoneinfo" could not be resolved (reportMissingImports)
+    import zoneinfo  # type: ignore[no-redef]
+
 # Until https://github.com/numpy/numpy/issues/19078 is sorted out, just suppress
 suppress_npdev_promotion_warning = pytest.mark.filterwarnings(
     "ignore:Promotion of numbers and bools:FutureWarning"
@@ -1166,6 +1171,8 @@ def iris(datapath):
     timezone(timedelta(hours=1)),
     timezone(timedelta(hours=-1), name="foo"),
 ]
+if zoneinfo is not None:
+    TIMEZONES.extend([zoneinfo.ZoneInfo("US/Pacific"), zoneinfo.ZoneInfo("UTC")])
 TIMEZONE_IDS = [repr(i) for i in TIMEZONES]
 
 
@@ -1191,7 +1198,12 @@ def tz_aware_fixture(request):
 tz_aware_fixture2 = tz_aware_fixture
 
 
-@pytest.fixture(params=["utc", "dateutil/UTC", utc, tzutc(), timezone.utc])
+_UTCS = ["utc", "dateutil/UTC", utc, tzutc(), timezone.utc]
+if zoneinfo is not None:
+    _UTCS.append(zoneinfo.ZoneInfo("UTC"))
+
+
+@pytest.fixture(params=_UTCS)
 def utc_fixture(request):
     """
     Fixture to provide variants of UTC timezone strings and tzinfo objects.

pandas/tests/indexes/datetimes/test_constructors.py

@@ -15,6 +15,7 @@
     OutOfBoundsDatetime,
     conversion,
 )
+from pandas.compat import PY39
 
 import pandas as pd
 from pandas import (
@@ -31,6 +32,9 @@
     period_array,
 )
 
+if PY39:
+    import zoneinfo
+
 
 class TestDatetimeIndex:
     @pytest.mark.parametrize(
@@ -1128,7 +1132,12 @@ def test_timestamp_constructor_retain_fold(tz, fold):
     assert result == expected
 
 
-@pytest.mark.parametrize("tz", ["dateutil/Europe/London"])
+_tzs = ["dateutil/Europe/London"]
+if PY39:
+    _tzs = ["dateutil/Europe/London", zoneinfo.ZoneInfo("Europe/London")]
+
+
+@pytest.mark.parametrize("tz", _tzs)
 @pytest.mark.parametrize(
     "ts_input,fold_out",
     [
@@ -1148,6 +1157,7 @@ def test_timestamp_constructor_infer_fold_from_value(tz, ts_input, fold_out):
     result = ts.fold
     expected = fold_out
     assert result == expected
+    # TODO: belongs in Timestamp tests?
 
 
 @pytest.mark.parametrize("tz", ["dateutil/Europe/London"])