Make DTA/TDA/PA return NotImplemented on comparisons (pandas-dev#24643)

Author: jbrockmendel

pandas/core/arrays/datetimes.py

@@ -19,7 +19,8 @@
     is_extension_type, is_float_dtype, is_object_dtype, is_period_dtype,
     is_string_dtype, is_timedelta64_dtype, pandas_dtype)
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
-from pandas.core.dtypes.generic import ABCIndexClass, ABCPandasArray, ABCSeries
+from pandas.core.dtypes.generic import (
+    ABCDataFrame, ABCIndexClass, ABCPandasArray, ABCSeries)
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
@@ -96,9 +97,8 @@ def _dt_array_cmp(cls, op):
     nat_result = True if opname == '__ne__' else False
 
     def wrapper(self, other):
-        # TODO: return NotImplemented for Series / Index and let pandas unbox
-        # Right now, returning NotImplemented for Index fails because we
-        # go into the index implementation, which may be a bug?
+        if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
+            return NotImplemented
 
         other = lib.item_from_zerodim(other)
 

pandas/core/arrays/period.py

@@ -17,7 +17,8 @@
     _TD_DTYPE, ensure_object, is_datetime64_dtype, is_float_dtype,
     is_list_like, is_period_dtype, pandas_dtype)
 from pandas.core.dtypes.dtypes import PeriodDtype
-from pandas.core.dtypes.generic import ABCIndexClass, ABCPeriodIndex, ABCSeries
+from pandas.core.dtypes.generic import (
+    ABCDataFrame, ABCIndexClass, ABCPeriodIndex, ABCSeries)
 from pandas.core.dtypes.missing import isna, notna
 
 import pandas.core.algorithms as algos
@@ -48,26 +49,20 @@ def _period_array_cmp(cls, op):
 
     def wrapper(self, other):
         op = getattr(self.asi8, opname)
-        # We want to eventually defer to the Series or PeriodIndex (which will
-        # return here with an unboxed PeriodArray). But before we do that,
-        # we do a bit of validation on type (Period) and freq, so that our
-        # error messages are sensible
+
+        if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
+            return NotImplemented
+
         if is_list_like(other) and len(other) != len(self):
             raise ValueError("Lengths must match")
 
-        not_implemented = isinstance(other, (ABCSeries, ABCIndexClass))
-        if not_implemented:
-            other = other._values
-
         if isinstance(other, Period):
             self._check_compatible_with(other)
 
             result = op(other.ordinal)
         elif isinstance(other, cls):
             self._check_compatible_with(other)
 
-            if not_implemented:
-                return NotImplemented
             result = op(other.asi8)
 
             mask = self._isnan | other._isnan

pandas/core/arrays/timedeltas.py

@@ -64,6 +64,9 @@ def _td_array_cmp(cls, op):
     nat_result = True if opname == '__ne__' else False
 
     def wrapper(self, other):
+        if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
+            return NotImplemented
+
         if _is_convertible_to_td(other) or other is NaT:
             try:
                 other = Timedelta(other)

pandas/core/indexes/datetimelike.py

@@ -109,6 +109,11 @@ def _create_comparison_method(cls, op):
         Create a comparison method that dispatches to ``cls.values``.
         """
         def wrapper(self, other):
+            if isinstance(other, ABCSeries):
+                # the arrays defer to Series for comparison ops but the indexes
+                #  don't, so we have to unwrap here.
+                other = other._values
+
             result = op(self._data, maybe_unwrap_index(other))
             return result
 

pandas/tests/arithmetic/test_period.py

@@ -152,7 +152,10 @@ def test_parr_cmp_pi_mismatched_freq_raises(self, freq, box_with_array):
 
         # TODO: Could parametrize over boxes for idx?
         idx = PeriodIndex(['2011', '2012', '2013', '2014'], freq='A')
-        with pytest.raises(IncompatibleFrequency, match=msg):
+        rev_msg = (r'Input has different freq=(M|2M|3M) from '
+                   r'PeriodArray\(freq=A-DEC\)')
+        idx_msg = rev_msg if box_with_array is tm.to_array else msg
+        with pytest.raises(IncompatibleFrequency, match=idx_msg):
             base <= idx
 
         # Different frequency
@@ -164,7 +167,10 @@ def test_parr_cmp_pi_mismatched_freq_raises(self, freq, box_with_array):
             Period('2011', freq='4M') >= base
 
         idx = PeriodIndex(['2011', '2012', '2013', '2014'], freq='4M')
-        with pytest.raises(IncompatibleFrequency, match=msg):
+        rev_msg = (r'Input has different freq=(M|2M|3M) from '
+                   r'PeriodArray\(freq=4M\)')
+        idx_msg = rev_msg if box_with_array is tm.to_array else msg
+        with pytest.raises(IncompatibleFrequency, match=idx_msg):
             base <= idx
 
     @pytest.mark.parametrize('freq', ['M', '2M', '3M'])