TST/REF: parametrize arithmetic tests, simplify parts of core.ops (#26799)

Author: jbrockmendel

pandas/core/ops.py

@@ -1077,7 +1077,7 @@ def fill_binop(left, right, fill_value):
     return left, right
 
 
-def mask_cmp_op(x, y, op, allowed_types):
+def mask_cmp_op(x, y, op):
     """
     Apply the function `op` to only non-null points in x and y.
 
@@ -1086,16 +1086,14 @@ def mask_cmp_op(x, y, op, allowed_types):
     x : array-like
     y : array-like
     op : binary operation
-    allowed_types : class or tuple of classes
 
     Returns
     -------
     result : ndarray[bool]
     """
-    # TODO: Can we make the allowed_types arg unnecessary?
     xrav = x.ravel()
     result = np.empty(x.size, dtype=bool)
-    if isinstance(y, allowed_types):
+    if isinstance(y, (np.ndarray, ABCSeries)):
         yrav = y.ravel()
         mask = notna(xrav) & notna(yrav)
         result[mask] = op(np.array(list(xrav[mask])),
@@ -1633,39 +1631,38 @@ def _arith_method_SERIES(cls, op, special):
                         if op in [divmod, rdivmod] else _construct_result)
 
     def na_op(x, y):
-        import pandas.core.computation.expressions as expressions
-        try:
-            result = expressions.evaluate(op, str_rep, x, y, **eval_kwargs)
-        except TypeError:
-            result = masked_arith_op(x, y, op)
-
-        result = missing.fill_zeros(result, x, y, op_name, fill_zeros)
-        return result
-
-    def safe_na_op(lvalues, rvalues):
         """
-        return the result of evaluating na_op on the passed in values
+        Return the result of evaluating op on the passed in values.
 
-        try coercion to object type if the native types are not compatible
+        If native types are not compatible, try coersion to object dtype.
 
         Parameters
         ----------
-        lvalues : array-like
-        rvalues : array-like
+        x : array-like
+        y : array-like or scalar
+
+        Returns
+        -------
+        array-like
 
         Raises
         ------
-        TypeError: invalid operation
+        TypeError : invalid operation
         """
+        import pandas.core.computation.expressions as expressions
         try:
-            with np.errstate(all='ignore'):
-                return na_op(lvalues, rvalues)
-        except Exception:
-            if is_object_dtype(lvalues):
-                return libalgos.arrmap_object(lvalues,
-                                              lambda x: op(x, rvalues))
+            result = expressions.evaluate(op, str_rep, x, y, **eval_kwargs)
+        except TypeError:
+            result = masked_arith_op(x, y, op)
+        except Exception:  # TODO: more specific?
+            if is_object_dtype(x):
+                return libalgos.arrmap_object(x,
+                                              lambda val: op(val, y))
             raise
 
+        result = missing.fill_zeros(result, x, y, op_name, fill_zeros)
+        return result
+
     def wrapper(left, right):
         if isinstance(right, ABCDataFrame):
             return NotImplemented
@@ -1713,7 +1710,8 @@ def wrapper(left, right):
         if isinstance(rvalues, ABCSeries):
             rvalues = rvalues.values
 
-        result = safe_na_op(lvalues, rvalues)
+        with np.errstate(all='ignore'):
+            result = na_op(lvalues, rvalues)
         return construct_result(left, result,
                                 index=left.index, name=res_name, dtype=None)
 
@@ -2136,7 +2134,6 @@ def na_op(x, y):
             result = masked_arith_op(x, y, op)
 
         result = missing.fill_zeros(result, x, y, op_name, fill_zeros)
-
         return result
 
     if op_name in _op_descriptions:
@@ -2183,7 +2180,7 @@ def na_op(x, y):
             with np.errstate(invalid='ignore'):
                 result = op(x, y)
         except TypeError:
-            result = mask_cmp_op(x, y, op, (np.ndarray, ABCSeries))
+            result = mask_cmp_op(x, y, op)
         return result
 
     doc = _flex_comp_doc_FRAME.format(op_name=op_name,

pandas/tests/arithmetic/test_datetime64.py

@@ -37,6 +37,27 @@ def assert_all(obj):
 # ------------------------------------------------------------------
 # Comparisons
 
+class TestDatetime64ArrayLikeComparisons:
+    # Comparison tests for datetime64 vectors fully parametrized over
+    #  DataFrame/Series/DatetimeIndex/DateteimeArray.  Ideally all comparison
+    #  tests will eventually end up here.
+
+    def test_compare_zerodim(self, tz_naive_fixture, box_with_array):
+        # Test comparison with zero-dimensional array is unboxed
+        tz = tz_naive_fixture
+        box = box_with_array
+        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+        dti = date_range('20130101', periods=3, tz=tz)
+
+        other = np.array(dti.to_numpy()[0])
+
+        # FIXME: ValueError with transpose on tzaware
+        dtarr = tm.box_expected(dti, box, transpose=False)
+        result = dtarr <= other
+        expected = np.array([True, False, False])
+        expected = tm.box_expected(expected, xbox, transpose=False)
+        tm.assert_equal(result, expected)
+
 
 class TestDatetime64DataFrameComparison:
     @pytest.mark.parametrize('timestamps', [
@@ -339,17 +360,6 @@ def test_comparison_tzawareness_compat(self, op):
 
 class TestDatetimeIndexComparisons:
 
-    # TODO: parametrize over box
-    def test_compare_zerodim(self, tz_naive_fixture):
-        # Test comparison with zero-dimensional array is unboxed
-        tz = tz_naive_fixture
-        dti = date_range('20130101', periods=3, tz=tz)
-
-        other = np.array(dti.to_numpy()[0])
-        result = dti <= other
-        expected = np.array([True, False, False])
-        tm.assert_numpy_array_equal(result, expected)
-
     # TODO: moved from tests.indexes.test_base; parametrize and de-duplicate
     @pytest.mark.parametrize("op", [
         operator.eq, operator.ne, operator.gt, operator.lt,

pandas/tests/arithmetic/test_period.py

@@ -20,17 +20,27 @@
 # Comparisons
 
 
-class TestPeriodIndexComparisons:
+class TestPeriodArrayLikeComparisons:
+    # Comparison tests for PeriodDtype vectors fully parametrized over
+    #  DataFrame/Series/PeriodIndex/PeriodArray.  Ideally all comparison
+    #  tests will eventually end up here.
 
-    # TODO: parameterize over boxes
-    def test_compare_zerodim(self):
+    def test_compare_zerodim(self, box_with_array):
         # GH#26689 make sure we unbox zero-dimensional arrays
+        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+
         pi = pd.period_range('2000', periods=4)
         other = np.array(pi.to_numpy()[0])
 
+        pi = tm.box_expected(pi, box_with_array)
         result = pi <= other
         expected = np.array([True, False, False, False])
-        tm.assert_numpy_array_equal(result, expected)
+        expected = tm.box_expected(expected, xbox)
+        tm.assert_equal(result, expected)
+
+
+class TestPeriodIndexComparisons:
+    # TODO: parameterize over boxes
 
     @pytest.mark.parametrize("other", ["2017", 2017])
     def test_eq(self, other):

pandas/tests/arithmetic/test_timedelta64.py

@@ -31,22 +31,33 @@ def get_upcast_box(box, vector):
 # ------------------------------------------------------------------
 # Timedelta64[ns] dtype Comparisons
 
-class TestTimedelta64ArrayComparisons:
-    # TODO: All of these need to be parametrized over box
+class TestTimedelta64ArrayLikeComparisons:
+    # Comparison tests for timedelta64[ns] vectors fully parametrized over
+    #  DataFrame/Series/TimedeltaIndex/TimedeltaArray.  Ideally all comparison
+    #  tests will eventually end up here.
 
-    def test_compare_timedelta64_zerodim(self):
+    def test_compare_timedelta64_zerodim(self, box_with_array):
         # GH#26689 should unbox when comparing with zerodim array
+        box = box_with_array
+        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+
         tdi = pd.timedelta_range('2H', periods=4)
         other = np.array(tdi.to_numpy()[0])
 
+        tdi = tm.box_expected(tdi, box)
         res = tdi <= other
         expected = np.array([True, False, False, False])
-        tm.assert_numpy_array_equal(res, expected)
+        expected = tm.box_expected(expected, xbox)
+        tm.assert_equal(res, expected)
 
         with pytest.raises(TypeError):
             # zero-dim of wrong dtype should still raise
             tdi >= np.array(4)
 
+
+class TestTimedelta64ArrayComparisons:
+    # TODO: All of these need to be parametrized over box
+
     def test_compare_timedelta_series(self):
         # regression test for GH#5963
         s = pd.Series([timedelta(days=1), timedelta(days=2)])

pandas/tests/tseries/offsets/test_offsets_properties.py

@@ -71,7 +71,10 @@ def test_on_offset_implementations(dt, offset):
     assert offset.onOffset(dt) == (compare == dt)
 
 
-@pytest.mark.xfail
+@pytest.mark.xfail(reason="res_v2 below is incorrect, needs to use the "
+                          "commented-out version with tz_localize.  "
+                          "But with that fix in place, hypothesis then "
+                          "has errors in timezone generation.")
 @given(gen_yqm_offset, gen_date_range)
 def test_apply_index_implementations(offset, rng):
     # offset.apply_index(dti)[i] should match dti[i] + offset
@@ -82,6 +85,7 @@ def test_apply_index_implementations(offset, rng):
 
     res = rng + offset
     res_v2 = offset.apply_index(rng)
+    # res_v2 = offset.apply_index(rng.tz_localize(None)).tz_localize(rng.tz)
     assert (res == res_v2).all()
 
     assert res[0] == rng[0] + offset
@@ -93,7 +97,7 @@ def test_apply_index_implementations(offset, rng):
     # TODO: Check randomly assorted entries, not just first/last
 
 
-@pytest.mark.xfail
+@pytest.mark.xfail  # TODO: reason?
 @given(gen_yqm_offset)
 def test_shift_across_dst(offset):
     # GH#18319 check that 1) timezone is correctly normalized and