TST: parametrize arithmetic tests (#27847)

Author: jbrockmendel

pandas/tests/arithmetic/test_datetime64.py

@@ -30,6 +30,54 @@
 import pandas.util.testing as tm
 
 
+def assert_invalid_comparison(left, right, box):
+    """
+    Assert that comparison operations with mismatched types behave correctly.
+
+    Parameters
+    ----------
+    left : np.ndarray, ExtensionArray, Index, or Series
+    right : object
+    box : {pd.DataFrame, pd.Series, pd.Index, tm.to_array}
+    """
+    # Not for tznaive-tzaware comparison
+
+    # Note: not quite the same as how we do this for tm.box_expected
+    xbox = box if box is not pd.Index else np.array
+
+    result = left == right
+    expected = xbox(np.zeros(result.shape, dtype=np.bool_))
+
+    tm.assert_equal(result, expected)
+
+    result = right == left
+    tm.assert_equal(result, expected)
+
+    result = left != right
+    tm.assert_equal(result, ~expected)
+
+    result = right != left
+    tm.assert_equal(result, ~expected)
+
+    msg = "Invalid comparison between"
+    with pytest.raises(TypeError, match=msg):
+        left < right
+    with pytest.raises(TypeError, match=msg):
+        left <= right
+    with pytest.raises(TypeError, match=msg):
+        left > right
+    with pytest.raises(TypeError, match=msg):
+        left >= right
+    with pytest.raises(TypeError, match=msg):
+        right < left
+    with pytest.raises(TypeError, match=msg):
+        right <= left
+    with pytest.raises(TypeError, match=msg):
+        right > left
+    with pytest.raises(TypeError, match=msg):
+        right >= left
+
+
 def assert_all(obj):
     """
     Test helper to call call obj.all() the appropriate number of times on
@@ -47,7 +95,7 @@ def assert_all(obj):
 
 class TestDatetime64ArrayLikeComparisons:
     # Comparison tests for datetime64 vectors fully parametrized over
-    #  DataFrame/Series/DatetimeIndex/DateteimeArray.  Ideally all comparison
+    #  DataFrame/Series/DatetimeIndex/DatetimeArray.  Ideally all comparison
     #  tests will eventually end up here.
 
     def test_compare_zerodim(self, tz_naive_fixture, box_with_array):
@@ -59,36 +107,61 @@ def test_compare_zerodim(self, tz_naive_fixture, box_with_array):
 
         other = np.array(dti.to_numpy()[0])
 
-        # FIXME: ValueError with transpose on tzaware
-        dtarr = tm.box_expected(dti, box, transpose=False)
+        dtarr = tm.box_expected(dti, box)
         result = dtarr <= other
         expected = np.array([True, False, False])
-        expected = tm.box_expected(expected, xbox, transpose=False)
+        expected = tm.box_expected(expected, xbox)
         tm.assert_equal(result, expected)
 
+    def test_dt64arr_cmp_date_invalid(self, tz_naive_fixture, box_with_array):
+        # GH#19800, GH#19301 datetime.date comparison raises to
+        #  match DatetimeIndex/Timestamp.  This also matches the behavior
+        #  of stdlib datetime.datetime
+        tz = tz_naive_fixture
 
-class TestDatetime64DataFrameComparison:
-    @pytest.mark.parametrize(
-        "timestamps",
-        [
-            [pd.Timestamp("2012-01-01 13:00:00+00:00")] * 2,
-            [pd.Timestamp("2012-01-01 13:00:00")] * 2,
-        ],
-    )
-    def test_tz_aware_scalar_comparison(self, timestamps):
-        # GH#15966
-        df = pd.DataFrame({"test": timestamps})
-        expected = pd.DataFrame({"test": [False, False]})
-        tm.assert_frame_equal(df == -1, expected)
+        dti = pd.date_range("20010101", periods=10, tz=tz)
+        date = dti[0].to_pydatetime().date()
+
+        dtarr = tm.box_expected(dti, box_with_array)
+        assert_invalid_comparison(dtarr, date, box_with_array)
 
-    def test_dt64_nat_comparison(self):
+    @pytest.mark.parametrize("other", ["foo", -1, 99, 4.0, object(), timedelta(days=2)])
+    def test_dt64arr_cmp_scalar_invalid(self, other, tz_naive_fixture, box_with_array):
+        # GH#22074, GH#15966
+        tz = tz_naive_fixture
+
+        rng = date_range("1/1/2000", periods=10, tz=tz)
+        dtarr = tm.box_expected(rng, box_with_array)
+        assert_invalid_comparison(dtarr, other, box_with_array)
+
+    @pytest.mark.parametrize("other", [None, np.nan])
+    def test_dt64arr_cmp_na_scalar_invalid(
+        self, other, tz_naive_fixture, box_with_array
+    ):
+        # GH#19301
+        tz = tz_naive_fixture
+        dti = pd.date_range("2016-01-01", periods=2, tz=tz)
+        dtarr = tm.box_expected(dti, box_with_array)
+        assert_invalid_comparison(dtarr, other, box_with_array)
+
+    def test_dt64arr_nat_comparison(self, tz_naive_fixture, box_with_array):
         # GH#22242, GH#22163 DataFrame considered NaT == ts incorrectly
-        ts = pd.Timestamp.now()
-        df = pd.DataFrame([ts, pd.NaT])
-        expected = pd.DataFrame([True, False])
+        tz = tz_naive_fixture
+        box = box_with_array
+        xbox = box if box is not pd.Index else np.ndarray
+
+        ts = pd.Timestamp.now(tz)
+        ser = pd.Series([ts, pd.NaT])
+
+        # FIXME: Can't transpose because that loses the tz dtype on
+        #  the NaT column
+        obj = tm.box_expected(ser, box, transpose=False)
 
-        result = df == ts
-        tm.assert_frame_equal(result, expected)
+        expected = pd.Series([True, False], dtype=np.bool_)
+        expected = tm.box_expected(expected, xbox, transpose=False)
+
+        result = obj == ts
+        tm.assert_equal(result, expected)
 
 
 class TestDatetime64SeriesComparison:
@@ -142,35 +215,17 @@ def test_nat_comparisons(self, dtype, box, reverse, pair):
         expected = Series([False, False, True])
         tm.assert_series_equal(left <= right, expected)
 
-    def test_comparison_invalid(self, box_with_array):
+    def test_comparison_invalid(self, tz_naive_fixture, box_with_array):
         # GH#4968
         # invalid date/int comparisons
-        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
-
+        tz = tz_naive_fixture
         ser = Series(range(5))
-        ser2 = Series(pd.date_range("20010101", periods=5))
+        ser2 = Series(pd.date_range("20010101", periods=5, tz=tz))
 
         ser = tm.box_expected(ser, box_with_array)
         ser2 = tm.box_expected(ser2, box_with_array)
 
-        for (x, y) in [(ser, ser2), (ser2, ser)]:
-
-            result = x == y
-            expected = tm.box_expected([False] * 5, xbox)
-            tm.assert_equal(result, expected)
-
-            result = x != y
-            expected = tm.box_expected([True] * 5, xbox)
-            tm.assert_equal(result, expected)
-            msg = "Invalid comparison between"
-            with pytest.raises(TypeError, match=msg):
-                x >= y
-            with pytest.raises(TypeError, match=msg):
-                x > y
-            with pytest.raises(TypeError, match=msg):
-                x < y
-            with pytest.raises(TypeError, match=msg):
-                x <= y
+        assert_invalid_comparison(ser, ser2, box_with_array)
 
     @pytest.mark.parametrize(
         "data",
@@ -227,26 +282,6 @@ def test_series_comparison_scalars(self):
         expected = Series([x > val for x in series])
         tm.assert_series_equal(result, expected)
 
-    def test_dt64ser_cmp_date_invalid(self, box_with_array):
-        # GH#19800 datetime.date comparison raises to
-        # match DatetimeIndex/Timestamp.  This also matches the behavior
-        # of stdlib datetime.datetime
-
-        ser = pd.date_range("20010101", periods=10)
-        date = ser[0].to_pydatetime().date()
-
-        ser = tm.box_expected(ser, box_with_array)
-        assert_all(~(ser == date))
-        assert_all(ser != date)
-        with pytest.raises(TypeError):
-            ser > date
-        with pytest.raises(TypeError):
-            ser < date
-        with pytest.raises(TypeError):
-            ser >= date
-        with pytest.raises(TypeError):
-            ser <= date
-
     @pytest.mark.parametrize(
         "left,right", [("lt", "gt"), ("le", "ge"), ("eq", "eq"), ("ne", "ne")]
     )
@@ -388,57 +423,6 @@ def test_dti_cmp_datetimelike(self, other, tz_naive_fixture):
         expected = np.array([True, False])
         tm.assert_numpy_array_equal(result, expected)
 
-    def dt64arr_cmp_non_datetime(self, tz_naive_fixture, box_with_array):
-        # GH#19301 by convention datetime.date is not considered comparable
-        # to Timestamp or DatetimeIndex.  This may change in the future.
-        tz = tz_naive_fixture
-        dti = pd.date_range("2016-01-01", periods=2, tz=tz)
-        dtarr = tm.box_expected(dti, box_with_array)
-
-        other = datetime(2016, 1, 1).date()
-        assert not (dtarr == other).any()
-        assert (dtarr != other).all()
-        with pytest.raises(TypeError):
-            dtarr < other
-        with pytest.raises(TypeError):
-            dtarr <= other
-        with pytest.raises(TypeError):
-            dtarr > other
-        with pytest.raises(TypeError):
-            dtarr >= other
-
-    @pytest.mark.parametrize("other", [None, np.nan, pd.NaT])
-    def test_dti_eq_null_scalar(self, other, tz_naive_fixture):
-        # GH#19301
-        tz = tz_naive_fixture
-        dti = pd.date_range("2016-01-01", periods=2, tz=tz)
-        assert not (dti == other).any()
-
-    @pytest.mark.parametrize("other", [None, np.nan, pd.NaT])
-    def test_dti_ne_null_scalar(self, other, tz_naive_fixture):
-        # GH#19301
-        tz = tz_naive_fixture
-        dti = pd.date_range("2016-01-01", periods=2, tz=tz)
-        assert (dti != other).all()
-
-    @pytest.mark.parametrize("other", [None, np.nan])
-    def test_dti_cmp_null_scalar_inequality(
-        self, tz_naive_fixture, other, box_with_array
-    ):
-        # GH#19301
-        tz = tz_naive_fixture
-        dti = pd.date_range("2016-01-01", periods=2, tz=tz)
-        dtarr = tm.box_expected(dti, box_with_array)
-        msg = "Invalid comparison between"
-        with pytest.raises(TypeError, match=msg):
-            dtarr < other
-        with pytest.raises(TypeError, match=msg):
-            dtarr <= other
-        with pytest.raises(TypeError, match=msg):
-            dtarr > other
-        with pytest.raises(TypeError, match=msg):
-            dtarr >= other
-
     @pytest.mark.parametrize("dtype", [None, object])
     def test_dti_cmp_nat(self, dtype, box_with_array):
         if box_with_array is tm.to_array and dtype is object:
@@ -728,34 +712,6 @@ def test_dti_cmp_str(self, tz_naive_fixture):
         expected = np.array([True] * 10)
         tm.assert_numpy_array_equal(result, expected)
 
-    @pytest.mark.parametrize("other", ["foo", 99, 4.0, object(), timedelta(days=2)])
-    def test_dt64arr_cmp_scalar_invalid(self, other, tz_naive_fixture, box_with_array):
-        # GH#22074
-        tz = tz_naive_fixture
-        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
-
-        rng = date_range("1/1/2000", periods=10, tz=tz)
-        rng = tm.box_expected(rng, box_with_array)
-
-        result = rng == other
-        expected = np.array([False] * 10)
-        expected = tm.box_expected(expected, xbox)
-        tm.assert_equal(result, expected)
-
-        result = rng != other
-        expected = np.array([True] * 10)
-        expected = tm.box_expected(expected, xbox)
-        tm.assert_equal(result, expected)
-        msg = "Invalid comparison between"
-        with pytest.raises(TypeError, match=msg):
-            rng < other
-        with pytest.raises(TypeError, match=msg):
-            rng <= other
-        with pytest.raises(TypeError, match=msg):
-            rng > other
-        with pytest.raises(TypeError, match=msg):
-            rng >= other
-
     def test_dti_cmp_list(self):
         rng = date_range("1/1/2000", periods=10)
 
@@ -2576,24 +2532,3 @@ def test_shift_months(years, months):
     raw = [x + pd.offsets.DateOffset(years=years, months=months) for x in dti]
     expected = DatetimeIndex(raw)
     tm.assert_index_equal(actual, expected)
-
-
-# FIXME: this belongs in scalar tests
-class SubDatetime(datetime):
-    pass
-
-
-@pytest.mark.parametrize(
-    "lh,rh",
-    [
-        (SubDatetime(2000, 1, 1), Timedelta(hours=1)),
-        (Timedelta(hours=1), SubDatetime(2000, 1, 1)),
-    ],
-)
-def test_dt_subclass_add_timedelta(lh, rh):
-    # GH 25851
-    # ensure that subclassed datetime works for
-    # Timedelta operations
-    result = lh + rh
-    expected = SubDatetime(2000, 1, 1, 1)
-    assert result == expected

pandas/tests/arithmetic/test_period.py

@@ -573,12 +573,19 @@ def test_parr_add_sub_float_raises(self, op, other, box_with_array):
     @pytest.mark.parametrize(
         "other",
         [
+            # datetime scalars
             pd.Timestamp.now(),
             pd.Timestamp.now().to_pydatetime(),
             pd.Timestamp.now().to_datetime64(),
+            # datetime-like arrays
+            pd.date_range("2016-01-01", periods=3, freq="H"),
+            pd.date_range("2016-01-01", periods=3, tz="Europe/Brussels"),
+            pd.date_range("2016-01-01", periods=3, freq="S")._data,
+            pd.date_range("2016-01-01", periods=3, tz="Asia/Tokyo")._data,
+            # Miscellaneous invalid types
         ],
     )
-    def test_parr_add_sub_datetime_scalar(self, other, box_with_array):
+    def test_parr_add_sub_invalid(self, other, box_with_array):
         # GH#23215
         rng = pd.period_range("1/1/2000", freq="D", periods=3)
         rng = tm.box_expected(rng, box_with_array)
@@ -595,23 +602,6 @@ def test_parr_add_sub_datetime_scalar(self, other, box_with_array):
     # -----------------------------------------------------------------
     # __add__/__sub__ with ndarray[datetime64] and ndarray[timedelta64]
 
-    def test_parr_add_sub_dt64_array_raises(self, box_with_array):
-        rng = pd.period_range("1/1/2000", freq="D", periods=3)
-        dti = pd.date_range("2016-01-01", periods=3)
-        dtarr = dti.values
-
-        rng = tm.box_expected(rng, box_with_array)
-
-        with pytest.raises(TypeError):
-            rng + dtarr
-        with pytest.raises(TypeError):
-            dtarr + rng
-
-        with pytest.raises(TypeError):
-            rng - dtarr
-        with pytest.raises(TypeError):
-            dtarr - rng
-
     def test_pi_add_sub_td64_array_non_tick_raises(self):
         rng = pd.period_range("1/1/2000", freq="Q", periods=3)
         tdi = pd.TimedeltaIndex(["-1 Day", "-1 Day", "-1 Day"])