test_tz_localize.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. from datetime import timezone
  2. import pytest
  3. import pytz
  4. from pandas._libs.tslibs import timezones
  5. from pandas import (
  6. DatetimeIndex,
  7. NaT,
  8. Series,
  9. Timestamp,
  10. date_range,
  11. )
  12. import pandas._testing as tm
  13. class TestTZLocalize:
  14. def test_series_tz_localize_ambiguous_bool(self):
  15. # make sure that we are correctly accepting bool values as ambiguous
  16. # GH#14402
  17. ts = Timestamp("2015-11-01 01:00:03")
  18. expected0 = Timestamp("2015-11-01 01:00:03-0500", tz="US/Central")
  19. expected1 = Timestamp("2015-11-01 01:00:03-0600", tz="US/Central")
  20. ser = Series([ts])
  21. expected0 = Series([expected0])
  22. expected1 = Series([expected1])
  23. with tm.external_error_raised(pytz.AmbiguousTimeError):
  24. ser.dt.tz_localize("US/Central")
  25. result = ser.dt.tz_localize("US/Central", ambiguous=True)
  26. tm.assert_series_equal(result, expected0)
  27. result = ser.dt.tz_localize("US/Central", ambiguous=[True])
  28. tm.assert_series_equal(result, expected0)
  29. result = ser.dt.tz_localize("US/Central", ambiguous=False)
  30. tm.assert_series_equal(result, expected1)
  31. result = ser.dt.tz_localize("US/Central", ambiguous=[False])
  32. tm.assert_series_equal(result, expected1)
  33. def test_series_tz_localize_matching_index(self):
  34. # Matching the index of the result with that of the original series
  35. # GH 43080
  36. dt_series = Series(
  37. date_range(start="2021-01-01T02:00:00", periods=5, freq="1D"),
  38. index=[2, 6, 7, 8, 11],
  39. dtype="category",
  40. )
  41. result = dt_series.dt.tz_localize("Europe/Berlin")
  42. expected = Series(
  43. date_range(
  44. start="2021-01-01T02:00:00", periods=5, freq="1D", tz="Europe/Berlin"
  45. ),
  46. index=[2, 6, 7, 8, 11],
  47. )
  48. tm.assert_series_equal(result, expected)
  49. @pytest.mark.parametrize(
  50. "method, exp",
  51. [
  52. ["shift_forward", "2015-03-29 03:00:00"],
  53. ["shift_backward", "2015-03-29 01:59:59.999999999"],
  54. ["NaT", NaT],
  55. ["raise", None],
  56. ["foo", "invalid"],
  57. ],
  58. )
  59. def test_tz_localize_nonexistent(self, warsaw, method, exp):
  60. # GH 8917
  61. tz = warsaw
  62. n = 60
  63. dti = date_range(start="2015-03-29 02:00:00", periods=n, freq="min")
  64. ser = Series(1, index=dti)
  65. df = ser.to_frame()
  66. if method == "raise":
  67. with tm.external_error_raised(pytz.NonExistentTimeError):
  68. dti.tz_localize(tz, nonexistent=method)
  69. with tm.external_error_raised(pytz.NonExistentTimeError):
  70. ser.tz_localize(tz, nonexistent=method)
  71. with tm.external_error_raised(pytz.NonExistentTimeError):
  72. df.tz_localize(tz, nonexistent=method)
  73. elif exp == "invalid":
  74. msg = (
  75. "The nonexistent argument must be one of "
  76. "'raise', 'NaT', 'shift_forward', 'shift_backward' "
  77. "or a timedelta object"
  78. )
  79. with pytest.raises(ValueError, match=msg):
  80. dti.tz_localize(tz, nonexistent=method)
  81. with pytest.raises(ValueError, match=msg):
  82. ser.tz_localize(tz, nonexistent=method)
  83. with pytest.raises(ValueError, match=msg):
  84. df.tz_localize(tz, nonexistent=method)
  85. else:
  86. result = ser.tz_localize(tz, nonexistent=method)
  87. expected = Series(1, index=DatetimeIndex([exp] * n, tz=tz))
  88. tm.assert_series_equal(result, expected)
  89. result = df.tz_localize(tz, nonexistent=method)
  90. expected = expected.to_frame()
  91. tm.assert_frame_equal(result, expected)
  92. res_index = dti.tz_localize(tz, nonexistent=method)
  93. tm.assert_index_equal(res_index, expected.index)
  94. @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
  95. def test_series_tz_localize_empty(self, tzstr):
  96. # GH#2248
  97. ser = Series(dtype=object)
  98. ser2 = ser.tz_localize("utc")
  99. assert ser2.index.tz == timezone.utc
  100. ser2 = ser.tz_localize(tzstr)
  101. timezones.tz_compare(ser2.index.tz, timezones.maybe_get_tz(tzstr))