From 6fd48c4c56ec7a786ce6df701691e3f0dc28a4d2 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Tue, 21 Jan 2020 14:51:36 -0800 Subject: [PATCH] require tzawareness compat for get_loc/get_value --- pandas/core/indexes/datetimes.py | 17 +++++++---- pandas/tests/series/indexing/test_datetime.py | 30 +++++++++++-------- pandas/tests/series/test_timezones.py | 8 ++++- 3 files changed, 36 insertions(+), 19 deletions(-) diff --git a/pandas/core/indexes/datetimes.py b/pandas/core/indexes/datetimes.py index ee9b948a76ac8..5198d5474b9a1 100644 --- a/pandas/core/indexes/datetimes.py +++ b/pandas/core/indexes/datetimes.py @@ -719,12 +719,19 @@ def get_loc(self, key, method=None, tolerance=None): def _maybe_cast_for_get_loc(self, key): # needed to localize naive datetimes - key = Timestamp(key) - if key.tzinfo is None: - key = key.tz_localize(self.tz) + stamp = Timestamp(key) + if not isinstance(key, str): + # strings get a pass on tzawareness compat, see GH#???? + try: + self._data._check_compatible_with(stamp) + except TypeError: + raise KeyError(key) + + if stamp.tzinfo is None: + stamp = stamp.tz_localize(self.tz) else: - key = key.tz_convert(self.tz) - return key + stamp = stamp.tz_convert(self.tz) + return stamp def _maybe_cast_slice_bound(self, label, side, kind): """ diff --git a/pandas/tests/series/indexing/test_datetime.py b/pandas/tests/series/indexing/test_datetime.py index 15ff5f6b343d1..da011d0e25989 100644 --- a/pandas/tests/series/indexing/test_datetime.py +++ b/pandas/tests/series/indexing/test_datetime.py @@ -1,4 +1,5 @@ from datetime import datetime, timedelta +import re import numpy as np import pytest @@ -289,24 +290,27 @@ def test_getitem_setitem_datetimeindex(): expected = ts[4:8] tm.assert_series_equal(result, expected) - # repeat all the above with naive datetimes - result = ts[datetime(1990, 1, 1, 4)] - expected = ts[4] - assert result == expected + # But we do not give datetimes a pass on tzawareness compat + # TODO: do the same with Timestamps and dt64 + msg = "Cannot compare tz-naive and tz-aware datetime-like objects" + naive = datetime(1990, 1, 1, 4) + with pytest.raises(KeyError, match=re.escape(repr(naive))): + ts[naive] result = ts.copy() - result[datetime(1990, 1, 1, 4)] = 0 - result[datetime(1990, 1, 1, 4)] = ts[4] - tm.assert_series_equal(result, ts) + with pytest.raises(TypeError, match=msg): + result[datetime(1990, 1, 1, 4)] = 0 + with pytest.raises(TypeError, match=msg): + result[datetime(1990, 1, 1, 4)] = ts[4] - result = ts[datetime(1990, 1, 1, 4) : datetime(1990, 1, 1, 7)] - expected = ts[4:8] - tm.assert_series_equal(result, expected) + with pytest.raises(TypeError, match=msg): + ts[datetime(1990, 1, 1, 4) : datetime(1990, 1, 1, 7)] result = ts.copy() - result[datetime(1990, 1, 1, 4) : datetime(1990, 1, 1, 7)] = 0 - result[datetime(1990, 1, 1, 4) : datetime(1990, 1, 1, 7)] = ts[4:8] - tm.assert_series_equal(result, ts) + with pytest.raises(TypeError, match=msg): + result[datetime(1990, 1, 1, 4) : datetime(1990, 1, 1, 7)] = 0 + with pytest.raises(TypeError, match=msg): + result[datetime(1990, 1, 1, 4) : datetime(1990, 1, 1, 7)] = ts[4:8] lb = datetime(1990, 1, 1, 4) rb = datetime(1990, 1, 1, 7) diff --git a/pandas/tests/series/test_timezones.py b/pandas/tests/series/test_timezones.py index a363f927d10a9..bcc3d495f3d57 100644 --- a/pandas/tests/series/test_timezones.py +++ b/pandas/tests/series/test_timezones.py @@ -339,7 +339,13 @@ def test_series_truncate_datetimeindex_tz(self): # GH 9243 idx = date_range("4/1/2005", "4/30/2005", freq="D", tz="US/Pacific") s = Series(range(len(idx)), index=idx) - result = s.truncate(datetime(2005, 4, 2), datetime(2005, 4, 4)) + msg = "Cannot compare tz-naive and tz-aware datetime-like objects" + with pytest.raises(TypeError, match=msg): + s.truncate(datetime(2005, 4, 2), datetime(2005, 4, 4)) + + lb = idx[1] + ub = idx[3] + result = s.truncate(lb.to_pydatetime(), ub.to_pydatetime()) expected = Series([1, 2, 3], index=idx[1:4]) tm.assert_series_equal(result, expected)