BUG: avoid segfault in tz_localize with ZoneInfo (#49936)

Author: jbrockmendel

Diff for: pandas/_libs/tslibs/tzconversion.pyx

@@ -233,6 +233,7 @@ timedelta-like}
         int64_t shift_delta = 0
         ndarray[int64_t] result_a, result_b, dst_hours
         int64_t[::1] result
+        bint is_zi = False
         bint infer_dst = False, is_dst = False, fill = False
         bint shift_forward = False, shift_backward = False
         bint fill_nonexist = False
@@ -304,6 +305,7 @@ timedelta-like}
     # Determine whether each date lies left of the DST transition (store in
     # result_a) or right of the DST transition (store in result_b)
     if is_zoneinfo(tz):
+        is_zi = True
         result_a, result_b =_get_utc_bounds_zoneinfo(
             vals, tz, creso=creso
         )
@@ -385,6 +387,11 @@ timedelta-like}
                     # nonexistent times
                     new_local = val - remaining_mins - 1
 
+                if is_zi:
+                    raise NotImplementedError(
+                        "nonexistent shifting is not implemented with ZoneInfo tzinfos"
+                    )
+
                 delta_idx = bisect_right_i8(info.tdata, new_local, info.ntrans)
 
                 delta_idx = delta_idx - delta_idx_offset

Diff for: pandas/conftest.py

@@ -1895,3 +1895,16 @@ def using_copy_on_write() -> bool:
     Fixture to check if Copy-on-Write is enabled.
     """
     return pd.options.mode.copy_on_write and pd.options.mode.data_manager == "block"
+
+
+warsaws = ["Europe/Warsaw", "dateutil/Europe/Warsaw"]
+if zoneinfo is not None:
+    warsaws.append(zoneinfo.ZoneInfo("Europe/Warsaw"))
+
+
+@pytest.fixture(params=warsaws)
+def warsaw(request):
+    """
+    tzinfo for Europe/Warsaw using pytz, dateutil, or zoneinfo.
+    """
+    return request.param

Diff for: pandas/tests/indexes/datetimes/test_constructors.py

@@ -872,10 +872,16 @@ def test_constructor_with_ambiguous_keyword_arg(self):
         result = date_range(end=end, periods=2, ambiguous=False)
         tm.assert_index_equal(result, expected)
 
-    def test_constructor_with_nonexistent_keyword_arg(self):
+    def test_constructor_with_nonexistent_keyword_arg(self, warsaw, request):
         # GH 35297
+        if type(warsaw).__name__ == "ZoneInfo":
+            mark = pytest.mark.xfail(
+                reason="nonexistent-shift not yet implemented for ZoneInfo",
+                raises=NotImplementedError,
+            )
+            request.node.add_marker(mark)
 
-        timezone = "Europe/Warsaw"
+        timezone = warsaw
 
         # nonexistent keyword in start
         start = Timestamp("2015-03-29 02:30:00").tz_localize(

Diff for: pandas/tests/indexes/datetimes/test_timezones.py

@@ -649,30 +649,6 @@ def test_dti_tz_localize_bdate_range(self):
         localized = dr.tz_localize(pytz.utc)
         tm.assert_index_equal(dr_utc, localized)
 
-    @pytest.mark.parametrize("tz", ["Europe/Warsaw", "dateutil/Europe/Warsaw"])
-    @pytest.mark.parametrize(
-        "method, exp", [["NaT", pd.NaT], ["raise", None], ["foo", "invalid"]]
-    )
-    def test_dti_tz_localize_nonexistent(self, tz, method, exp):
-        # GH 8917
-        n = 60
-        dti = date_range(start="2015-03-29 02:00:00", periods=n, freq="min")
-        if method == "raise":
-            with pytest.raises(pytz.NonExistentTimeError, match="2015-03-29 02:00:00"):
-                dti.tz_localize(tz, nonexistent=method)
-        elif exp == "invalid":
-            msg = (
-                "The nonexistent argument must be one of "
-                "'raise', 'NaT', 'shift_forward', 'shift_backward' "
-                "or a timedelta object"
-            )
-            with pytest.raises(ValueError, match=msg):
-                dti.tz_localize(tz, nonexistent=method)
-        else:
-            result = dti.tz_localize(tz, nonexistent=method)
-            expected = DatetimeIndex([exp] * n, tz=tz)
-            tm.assert_index_equal(result, expected)
-
     @pytest.mark.parametrize(
         "start_ts, tz, end_ts, shift",
         [
@@ -730,10 +706,9 @@ def test_dti_tz_localize_nonexistent_shift(
         tm.assert_index_equal(result, expected)
 
     @pytest.mark.parametrize("offset", [-1, 1])
-    @pytest.mark.parametrize("tz_type", ["", "dateutil/"])
-    def test_dti_tz_localize_nonexistent_shift_invalid(self, offset, tz_type):
+    def test_dti_tz_localize_nonexistent_shift_invalid(self, offset, warsaw):
         # GH 8917
-        tz = tz_type + "Europe/Warsaw"
+        tz = warsaw
         dti = DatetimeIndex([Timestamp("2015-03-29 02:20:00")])
         msg = "The provided timedelta will relocalize on a nonexistent time"
         with pytest.raises(ValueError, match=msg):

Diff for: pandas/tests/scalar/timestamp/test_timezones.py

@@ -141,9 +141,9 @@ def test_tz_localize_ambiguous_raise(self):
         with pytest.raises(AmbiguousTimeError, match=msg):
             ts.tz_localize("US/Pacific", ambiguous="raise")
 
-    def test_tz_localize_nonexistent_invalid_arg(self):
+    def test_tz_localize_nonexistent_invalid_arg(self, warsaw):
         # GH 22644
-        tz = "Europe/Warsaw"
+        tz = warsaw
         ts = Timestamp("2015-03-29 02:00:00")
         msg = (
             "The nonexistent argument must be one of 'raise', 'NaT', "
@@ -291,27 +291,26 @@ def test_timestamp_tz_localize_nonexistent_shift(
         assert result._creso == getattr(NpyDatetimeUnit, f"NPY_FR_{unit}").value
 
     @pytest.mark.parametrize("offset", [-1, 1])
-    @pytest.mark.parametrize("tz_type", ["", "dateutil/"])
-    def test_timestamp_tz_localize_nonexistent_shift_invalid(self, offset, tz_type):
+    def test_timestamp_tz_localize_nonexistent_shift_invalid(self, offset, warsaw):
         # GH 8917, 24466
-        tz = tz_type + "Europe/Warsaw"
+        tz = warsaw
         ts = Timestamp("2015-03-29 02:20:00")
         msg = "The provided timedelta will relocalize on a nonexistent time"
         with pytest.raises(ValueError, match=msg):
             ts.tz_localize(tz, nonexistent=timedelta(seconds=offset))
 
-    @pytest.mark.parametrize("tz", ["Europe/Warsaw", "dateutil/Europe/Warsaw"])
     @pytest.mark.parametrize("unit", ["ns", "us", "ms", "s"])
-    def test_timestamp_tz_localize_nonexistent_NaT(self, tz, unit):
+    def test_timestamp_tz_localize_nonexistent_NaT(self, warsaw, unit):
         # GH 8917
+        tz = warsaw
         ts = Timestamp("2015-03-29 02:20:00").as_unit(unit)
         result = ts.tz_localize(tz, nonexistent="NaT")
         assert result is NaT
 
-    @pytest.mark.parametrize("tz", ["Europe/Warsaw", "dateutil/Europe/Warsaw"])
     @pytest.mark.parametrize("unit", ["ns", "us", "ms", "s"])
-    def test_timestamp_tz_localize_nonexistent_raise(self, tz, unit):
+    def test_timestamp_tz_localize_nonexistent_raise(self, warsaw, unit):
         # GH 8917
+        tz = warsaw
         ts = Timestamp("2015-03-29 02:20:00").as_unit(unit)
         msg = "2015-03-29 02:20:00"
         with pytest.raises(pytz.NonExistentTimeError, match=msg):

Diff for: pandas/tests/series/methods/test_tz_localize.py

@@ -58,7 +58,6 @@ def test_series_tz_localize_matching_index(self):
         )
         tm.assert_series_equal(result, expected)
 
-    @pytest.mark.parametrize("tz", ["Europe/Warsaw", "dateutil/Europe/Warsaw"])
     @pytest.mark.parametrize(
         "method, exp",
         [
@@ -68,8 +67,9 @@ def test_series_tz_localize_matching_index(self):
             ["foo", "invalid"],
         ],
     )
-    def test_tz_localize_nonexistent(self, tz, method, exp):
+    def test_tz_localize_nonexistent(self, warsaw, method, exp):
         # GH 8917
+        tz = warsaw
         n = 60
         dti = date_range(start="2015-03-29 02:00:00", periods=n, freq="min")
         ser = Series(1, index=dti)
@@ -85,13 +85,27 @@ def test_tz_localize_nonexistent(self, tz, method, exp):
                 df.tz_localize(tz, nonexistent=method)
 
         elif exp == "invalid":
-            with pytest.raises(ValueError, match="argument must be one of"):
+            msg = (
+                "The nonexistent argument must be one of "
+                "'raise', 'NaT', 'shift_forward', 'shift_backward' "
+                "or a timedelta object"
+            )
+            with pytest.raises(ValueError, match=msg):
                 dti.tz_localize(tz, nonexistent=method)
-            with pytest.raises(ValueError, match="argument must be one of"):
+            with pytest.raises(ValueError, match=msg):
                 ser.tz_localize(tz, nonexistent=method)
-            with pytest.raises(ValueError, match="argument must be one of"):
+            with pytest.raises(ValueError, match=msg):
                 df.tz_localize(tz, nonexistent=method)
 
+        elif method == "shift_forward" and type(tz).__name__ == "ZoneInfo":
+            msg = "nonexistent shifting is not implemented with ZoneInfo tzinfos"
+            with pytest.raises(NotImplementedError, match=msg):
+                ser.tz_localize(tz, nonexistent=method)
+            with pytest.raises(NotImplementedError, match=msg):
+                df.tz_localize(tz, nonexistent=method)
+            with pytest.raises(NotImplementedError, match=msg):
+                dti.tz_localize(tz, nonexistent=method)
+
         else:
             result = ser.tz_localize(tz, nonexistent=method)
             expected = Series(1, index=DatetimeIndex([exp] * n, tz=tz))
@@ -101,6 +115,9 @@ def test_tz_localize_nonexistent(self, tz, method, exp):
             expected = expected.to_frame()
             tm.assert_frame_equal(result, expected)
 
+            res_index = dti.tz_localize(tz, nonexistent=method)
+            tm.assert_index_equal(res_index, expected.index)
+
     @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
     def test_series_tz_localize_empty(self, tzstr):
         # GH#2248