implement assert_tzawareness_compat for DatetimeIndex (#18376)

Author: jbrockmendel

doc/source/whatsnew/v0.23.0.txt

@@ -342,6 +342,7 @@ Conversion
 - Bug in :class:`Series`` with ``dtype='timedelta64[ns]`` where addition or subtraction of ``TimedeltaIndex`` had results cast to ``dtype='int64'`` (:issue:`17250`)
 - Bug in :class:`TimedeltaIndex` where division by a ``Series`` would return a ``TimedeltaIndex`` instead of a ``Series`` (issue:`19042`)
 - Bug in :class:`Series` with ``dtype='timedelta64[ns]`` where addition or subtraction of ``TimedeltaIndex`` could return a ``Series`` with an incorrect name (issue:`19043`)
+- Fixed bug where comparing :class:`DatetimeIndex` failed to raise ``TypeError`` when attempting to compare timezone-aware and timezone-naive datetimelike objects (:issue:`18162`)
 -
 
 Indexing

pandas/core/indexes/datetimes.py

@@ -13,14 +13,14 @@
     _INT64_DTYPE,
     _NS_DTYPE,
     is_object_dtype,
-    is_datetime64_dtype,
+    is_datetime64_dtype, is_datetime64tz_dtype,
     is_datetimetz,
     is_dtype_equal,
     is_timedelta64_dtype,
     is_integer,
     is_float,
     is_integer_dtype,
-    is_datetime64_ns_dtype,
+    is_datetime64_ns_dtype, is_datetimelike,
     is_period_dtype,
     is_bool_dtype,
     is_string_like,
@@ -106,8 +106,12 @@ def _dt_index_cmp(opname, cls, nat_result=False):
 
     def wrapper(self, other):
         func = getattr(super(DatetimeIndex, self), opname)
-        if (isinstance(other, datetime) or
-                isinstance(other, compat.string_types)):
+
+        if isinstance(other, (datetime, compat.string_types)):
+            if isinstance(other, datetime):
+                # GH#18435 strings get a pass from tzawareness compat
+                self._assert_tzawareness_compat(other)
+
             other = _to_m8(other, tz=self.tz)
             result = func(other)
             if isna(other):
@@ -117,6 +121,10 @@ def wrapper(self, other):
                 other = DatetimeIndex(other)
             elif not isinstance(other, (np.ndarray, Index, ABCSeries)):
                 other = _ensure_datetime64(other)
+
+            if is_datetimelike(other):
+                self._assert_tzawareness_compat(other)
+
             result = func(np.asarray(other))
             result = _values_from_object(result)
 
@@ -652,6 +660,20 @@ def _simple_new(cls, values, name=None, freq=None, tz=None,
         result._reset_identity()
         return result
 
+    def _assert_tzawareness_compat(self, other):
+        # adapted from _Timestamp._assert_tzawareness_compat
+        other_tz = getattr(other, 'tzinfo', None)
+        if is_datetime64tz_dtype(other):
+            # Get tzinfo from Series dtype
+            other_tz = other.dtype.tz
+        if self.tz is None:
+            if other_tz is not None:
+                raise TypeError('Cannot compare tz-naive and tz-aware '
+                                'datetime-like objects.')
+        elif other_tz is None:
+            raise TypeError('Cannot compare tz-naive and tz-aware '
+                            'datetime-like objects')
+
     @property
     def tzinfo(self):
         """

pandas/tests/indexes/datetimes/test_datetime.py

@@ -1,3 +1,5 @@
+import operator
+
 import pytest
 
 import numpy as np
@@ -248,6 +250,42 @@ def test_append_join_nondatetimeindex(self):
         # it works
         rng.join(idx, how='outer')
 
+    @pytest.mark.parametrize('op', [operator.eq, operator.ne,
+                                    operator.gt, operator.ge,
+                                    operator.lt, operator.le])
+    def test_comparison_tzawareness_compat(self, op):
+        # GH#18162
+        dr = pd.date_range('2016-01-01', periods=6)
+        dz = dr.tz_localize('US/Pacific')
+
+        with pytest.raises(TypeError):
+            op(dr, dz)
+        with pytest.raises(TypeError):
+            op(dr, list(dz))
+        with pytest.raises(TypeError):
+            op(dz, dr)
+        with pytest.raises(TypeError):
+            op(dz, list(dr))
+
+        # Check that there isn't a problem aware-aware and naive-naive do not
+        # raise
+        assert (dr == dr).all()
+        assert (dr == list(dr)).all()
+        assert (dz == dz).all()
+        assert (dz == list(dz)).all()
+
+        # Check comparisons against scalar Timestamps
+        ts = pd.Timestamp('2000-03-14 01:59')
+        ts_tz = pd.Timestamp('2000-03-14 01:59', tz='Europe/Amsterdam')
+
+        assert (dr > ts).all()
+        with pytest.raises(TypeError):
+            op(dr, ts_tz)
+
+        assert (dz > ts_tz).all()
+        with pytest.raises(TypeError):
+            op(dz, ts)
+
     def test_comparisons_coverage(self):
         rng = date_range('1/1/2000', periods=10)
 

pandas/tests/indexes/test_base.py

@@ -2262,6 +2262,26 @@ def test_intersect_str_dates(self):
 
         assert len(res) == 0
 
+    @pytest.mark.parametrize('op', [operator.eq, operator.ne,
+                                    operator.gt, operator.ge,
+                                    operator.lt, operator.le])
+    def test_comparison_tzawareness_compat(self, op):
+        # GH#18162
+        dr = pd.date_range('2016-01-01', periods=6)
+        dz = dr.tz_localize('US/Pacific')
+
+        # Check that there isn't a problem aware-aware and naive-naive do not
+        # raise
+        naive_series = Series(dr)
+        aware_series = Series(dz)
+        with pytest.raises(TypeError):
+            op(dz, naive_series)
+        with pytest.raises(TypeError):
+            op(dr, aware_series)
+
+        # TODO: implement _assert_tzawareness_compat for the reverse
+        # comparison with the Series on the left-hand side
+
 
 class TestIndexUtils(object):
 

pandas/tests/indexing/test_coercion.py

@@ -821,6 +821,9 @@ def test_replace_series(self, how, to_key, from_key):
         if (from_key.startswith('datetime') and to_key.startswith('datetime')):
             # tested below
             return
+        elif from_key in ['datetime64[ns, US/Eastern]', 'datetime64[ns, UTC]']:
+            # tested below
+            return
 
         if how == 'dict':
             replacer = dict(zip(self.rep[from_key], self.rep[to_key]))
@@ -849,6 +852,37 @@ def test_replace_series(self, how, to_key, from_key):
 
         tm.assert_series_equal(result, exp)
 
+    # TODO(jbrockmendel) commented out to only have a single xfail printed
+    @pytest.mark.xfail(reason='GH #18376, tzawareness-compat bug '
+                              'in BlockManager.replace_list')
+    # @pytest.mark.parametrize('how', ['dict', 'series'])
+    # @pytest.mark.parametrize('to_key', ['timedelta64[ns]', 'bool', 'object',
+    #                                     'complex128', 'float64', 'int64'])
+    # @pytest.mark.parametrize('from_key', ['datetime64[ns, UTC]',
+    #                                       'datetime64[ns, US/Eastern]'])
+    # def test_replace_series_datetime_tz(self, how, to_key, from_key):
+    def test_replace_series_datetime_tz(self):
+        how = 'series'
+        from_key = 'datetime64[ns, US/Eastern]'
+        to_key = 'timedelta64[ns]'
+
+        index = pd.Index([3, 4], name='xxx')
+        obj = pd.Series(self.rep[from_key], index=index, name='yyy')
+        assert obj.dtype == from_key
+
+        if how == 'dict':
+            replacer = dict(zip(self.rep[from_key], self.rep[to_key]))
+        elif how == 'series':
+            replacer = pd.Series(self.rep[to_key], index=self.rep[from_key])
+        else:
+            raise ValueError
+
+        result = obj.replace(replacer)
+        exp = pd.Series(self.rep[to_key], index=index, name='yyy')
+        assert exp.dtype == to_key
+
+        tm.assert_series_equal(result, exp)
+
     # TODO(jreback) commented out to only have a single xfail printed
     @pytest.mark.xfail(reason="different tz, "
                        "currently mask_missing raises SystemError")

pandas/tests/series/test_indexing.py

@@ -450,6 +450,13 @@ def test_getitem_setitem_datetimeindex(self):
 
         lb = "1990-01-01 04:00:00"
         rb = "1990-01-01 07:00:00"
+        # GH#18435 strings get a pass from tzawareness compat
+        result = ts[(ts.index >= lb) & (ts.index <= rb)]
+        expected = ts[4:8]
+        assert_series_equal(result, expected)
+
+        lb = "1990-01-01 04:00:00-0500"
+        rb = "1990-01-01 07:00:00-0500"
         result = ts[(ts.index >= lb) & (ts.index <= rb)]
         expected = ts[4:8]
         assert_series_equal(result, expected)
@@ -475,6 +482,13 @@ def test_getitem_setitem_datetimeindex(self):
 
         lb = datetime(1990, 1, 1, 4)
         rb = datetime(1990, 1, 1, 7)
+        with pytest.raises(TypeError):
+            # tznaive vs tzaware comparison is invalid
+            # see GH#18376, GH#18162
+            ts[(ts.index >= lb) & (ts.index <= rb)]
+
+        lb = pd.Timestamp(datetime(1990, 1, 1, 4)).tz_localize(rng.tzinfo)
+        rb = pd.Timestamp(datetime(1990, 1, 1, 7)).tz_localize(rng.tzinfo)
         result = ts[(ts.index >= lb) & (ts.index <= rb)]
         expected = ts[4:8]
         assert_series_equal(result, expected)

pandas/tests/test_base.py

@@ -114,7 +114,7 @@ def __init__(self, obj):
     def setup_method(self, method):
         pass
 
-    def test_invalida_delgation(self):
+    def test_invalid_delegation(self):
         # these show that in order for the delegation to work
         # the _delegate_* methods need to be overridden to not raise
         # a TypeError