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

Author: jbrockmendel

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

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

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(

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):

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):

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