1import pytest
2import pytz
3
4from pandas._libs.tslibs import timezones
5
6from pandas import DatetimeIndex, NaT, Series, Timestamp, date_range
7import pandas._testing as tm
8
9
10class TestTZLocalize:
11    def test_series_tz_localize(self):
12
13        rng = date_range("1/1/2011", periods=100, freq="H")
14        ts = Series(1, index=rng)
15
16        result = ts.tz_localize("utc")
17        assert result.index.tz.zone == "UTC"
18
19        # Can't localize if already tz-aware
20        rng = date_range("1/1/2011", periods=100, freq="H", tz="utc")
21        ts = Series(1, index=rng)
22
23        with pytest.raises(TypeError, match="Already tz-aware"):
24            ts.tz_localize("US/Eastern")
25
26    def test_series_tz_localize_ambiguous_bool(self):
27        # make sure that we are correctly accepting bool values as ambiguous
28
29        # GH#14402
30        ts = Timestamp("2015-11-01 01:00:03")
31        expected0 = Timestamp("2015-11-01 01:00:03-0500", tz="US/Central")
32        expected1 = Timestamp("2015-11-01 01:00:03-0600", tz="US/Central")
33
34        ser = Series([ts])
35        expected0 = Series([expected0])
36        expected1 = Series([expected1])
37
38        with tm.external_error_raised(pytz.AmbiguousTimeError):
39            ser.dt.tz_localize("US/Central")
40
41        result = ser.dt.tz_localize("US/Central", ambiguous=True)
42        tm.assert_series_equal(result, expected0)
43
44        result = ser.dt.tz_localize("US/Central", ambiguous=[True])
45        tm.assert_series_equal(result, expected0)
46
47        result = ser.dt.tz_localize("US/Central", ambiguous=False)
48        tm.assert_series_equal(result, expected1)
49
50        result = ser.dt.tz_localize("US/Central", ambiguous=[False])
51        tm.assert_series_equal(result, expected1)
52
53    @pytest.mark.parametrize("tz", ["Europe/Warsaw", "dateutil/Europe/Warsaw"])
54    @pytest.mark.parametrize(
55        "method, exp",
56        [
57            ["shift_forward", "2015-03-29 03:00:00"],
58            ["NaT", NaT],
59            ["raise", None],
60            ["foo", "invalid"],
61        ],
62    )
63    def test_series_tz_localize_nonexistent(self, tz, method, exp):
64        # GH 8917
65        n = 60
66        dti = date_range(start="2015-03-29 02:00:00", periods=n, freq="min")
67        s = Series(1, dti)
68        if method == "raise":
69            with tm.external_error_raised(pytz.NonExistentTimeError):
70                s.tz_localize(tz, nonexistent=method)
71        elif exp == "invalid":
72            with pytest.raises(ValueError, match="argument must be one of"):
73                dti.tz_localize(tz, nonexistent=method)
74        else:
75            result = s.tz_localize(tz, nonexistent=method)
76            expected = Series(1, index=DatetimeIndex([exp] * n, tz=tz))
77            tm.assert_series_equal(result, expected)
78
79    @pytest.mark.parametrize("tzstr", ["US/Eastern", "dateutil/US/Eastern"])
80    def test_series_tz_localize_empty(self, tzstr):
81        # GH#2248
82        ser = Series(dtype=object)
83
84        ser2 = ser.tz_localize("utc")
85        assert ser2.index.tz == pytz.utc
86
87        ser2 = ser.tz_localize(tzstr)
88        timezones.tz_compare(ser2.index.tz, timezones.maybe_get_tz(tzstr))
89