REF/TST: Collect Straggler Arithmetic Tests (pandas-dev#24376)

Author: jbrockmendel

pandas/tests/arithmetic/test_datetime64.py

@@ -20,6 +20,8 @@
 from pandas._libs.tslibs.conversion import localize_pydatetime
 from pandas._libs.tslibs.offsets import shift_months
 
+from pandas.core.indexes.datetimes import _to_m8
+
 from pandas import (
     Timestamp, Timedelta, Period, Series, date_range, NaT,
     DatetimeIndex, TimedeltaIndex)
@@ -61,6 +63,50 @@ def test_dt64_nat_comparison(self):
 
 class TestDatetime64SeriesComparison(object):
     # TODO: moved from tests.series.test_operators; needs cleanup
+
+    @pytest.mark.parametrize('pair', [
+        ([pd.Timestamp('2011-01-01'), NaT, pd.Timestamp('2011-01-03')],
+         [NaT, NaT, pd.Timestamp('2011-01-03')]),
+
+        ([pd.Timedelta('1 days'), NaT, pd.Timedelta('3 days')],
+         [NaT, NaT, pd.Timedelta('3 days')]),
+
+        ([pd.Period('2011-01', freq='M'), NaT,
+          pd.Period('2011-03', freq='M')],
+         [NaT, NaT, pd.Period('2011-03', freq='M')]),
+
+    ])
+    @pytest.mark.parametrize('reverse', [True, False])
+    @pytest.mark.parametrize('box', [Series, pd.Index])
+    @pytest.mark.parametrize('dtype', [None, object])
+    def test_nat_comparisons(self, dtype, box, reverse, pair):
+        l, r = pair
+        if reverse:
+            # add lhs / rhs switched data
+            l, r = r, l
+
+        left = Series(l, dtype=dtype)
+        right = box(r, dtype=dtype)
+        # Series, Index
+
+        expected = Series([False, False, True])
+        tm.assert_series_equal(left == right, expected)
+
+        expected = Series([True, True, False])
+        tm.assert_series_equal(left != right, expected)
+
+        expected = Series([False, False, False])
+        tm.assert_series_equal(left < right, expected)
+
+        expected = Series([False, False, False])
+        tm.assert_series_equal(left > right, expected)
+
+        expected = Series([False, False, True])
+        tm.assert_series_equal(left >= right, expected)
+
+        expected = Series([False, False, True])
+        tm.assert_series_equal(left <= right, expected)
+
     def test_comparison_invalid(self, box_with_array):
         # GH#4968
         # invalid date/int comparisons
@@ -272,8 +318,46 @@ def test_dt64arr_timestamp_equality(self, box_with_array):
         expected = tm.box_expected([False, False], xbox)
         tm.assert_equal(result, expected)
 
+    @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 TestDatetimeIndexComparisons(object):
+
+    # TODO: moved from tests.indexes.test_base; parametrize and de-duplicate
+    @pytest.mark.parametrize("op", [
+        operator.eq, operator.ne, operator.gt, operator.lt,
+        operator.ge, operator.le
+    ])
+    def test_comparators(self, op):
+        index = tm.makeDateIndex(100)
+        element = index[len(index) // 2]
+        element = _to_m8(element)
+
+        arr = np.array(index)
+        arr_result = op(arr, element)
+        index_result = op(index, element)
+
+        assert isinstance(index_result, np.ndarray)
+        tm.assert_numpy_array_equal(arr_result, index_result)
+
     @pytest.mark.parametrize('other', [datetime(2016, 1, 1),
                                        Timestamp('2016-01-01'),
                                        np.datetime64('2016-01-01')])
@@ -1450,6 +1534,42 @@ def test_datetimeindex_sub_datetimeindex_overflow(self):
 
 class TestTimestampSeriesArithmetic(object):
 
+    def test_empty_series_add_sub(self):
+        # GH#13844
+        a = Series(dtype='M8[ns]')
+        b = Series(dtype='m8[ns]')
+        tm.assert_series_equal(a, a + b)
+        tm.assert_series_equal(a, a - b)
+        tm.assert_series_equal(a, b + a)
+        with pytest.raises(TypeError):
+            b - a
+
+    def test_operators_datetimelike(self):
+
+        # ## timedelta64 ###
+        td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
+        td1.iloc[2] = np.nan
+
+        # ## datetime64 ###
+        dt1 = Series([pd.Timestamp('20111230'), pd.Timestamp('20120101'),
+                      pd.Timestamp('20120103')])
+        dt1.iloc[2] = np.nan
+        dt2 = Series([pd.Timestamp('20111231'), pd.Timestamp('20120102'),
+                      pd.Timestamp('20120104')])
+        dt1 - dt2
+        dt2 - dt1
+
+        # ## datetime64 with timetimedelta ###
+        dt1 + td1
+        td1 + dt1
+        dt1 - td1
+        # TODO: Decide if this ought to work.
+        # td1 - dt1
+
+        # ## timetimedelta with datetime64 ###
+        td1 + dt1
+        dt1 + td1
+
     def test_dt64ser_sub_datetime_dtype(self):
         ts = Timestamp(datetime(1993, 1, 7, 13, 30, 00))
         dt = datetime(1993, 6, 22, 13, 30)

pandas/tests/arithmetic/test_numeric.py

@@ -3,6 +3,7 @@
 # behave identically.
 # Specifically for numeric dtypes
 from decimal import Decimal
+from itertools import combinations
 import operator
 
 import pytest
@@ -846,11 +847,15 @@ def check(series, other):
 class TestUFuncCompat(object):
 
     @pytest.mark.parametrize('holder', [pd.Int64Index, pd.UInt64Index,
-                                        pd.Float64Index, pd.Series])
+                                        pd.Float64Index, pd.RangeIndex,
+                                        pd.Series])
     def test_ufunc_compat(self, holder):
         box = pd.Series if holder is pd.Series else pd.Index
 
-        idx = holder(np.arange(5, dtype='int64'))
+        if holder is pd.RangeIndex:
+            idx = pd.RangeIndex(0, 5)
+        else:
+            idx = holder(np.arange(5, dtype='int64'))
         result = np.sin(idx)
         expected = box(np.sin(np.arange(5, dtype='int64')))
         tm.assert_equal(result, expected)
@@ -940,3 +945,117 @@ def test_operators_reverse_object(self, op):
         result = op(1., arr)
         expected = op(1., arr.astype(float))
         tm.assert_series_equal(result.astype(float), expected)
+
+
+class TestNumericArithmeticUnsorted(object):
+    # Tests in this class have been moved from type-specific test modules
+    #  but not yet sorted, parametrized, and de-duplicated
+
+    def check_binop(self, ops, scalars, idxs):
+        for op in ops:
+            for a, b in combinations(idxs, 2):
+                result = op(a, b)
+                expected = op(pd.Int64Index(a), pd.Int64Index(b))
+                tm.assert_index_equal(result, expected)
+            for idx in idxs:
+                for scalar in scalars:
+                    result = op(idx, scalar)
+                    expected = op(pd.Int64Index(idx), scalar)
+                    tm.assert_index_equal(result, expected)
+
+    def test_binops(self):
+        ops = [operator.add, operator.sub, operator.mul, operator.floordiv,
+               operator.truediv]
+        scalars = [-1, 1, 2]
+        idxs = [pd.RangeIndex(0, 10, 1), pd.RangeIndex(0, 20, 2),
+                pd.RangeIndex(-10, 10, 2), pd.RangeIndex(5, -5, -1)]
+        self.check_binop(ops, scalars, idxs)
+
+    def test_binops_pow(self):
+        # later versions of numpy don't allow powers of negative integers
+        # so test separately
+        # https://github.com/numpy/numpy/pull/8127
+        ops = [pow]
+        scalars = [1, 2]
+        idxs = [pd.RangeIndex(0, 10, 1), pd.RangeIndex(0, 20, 2)]
+        self.check_binop(ops, scalars, idxs)
+
+    # TODO: mod, divmod?
+    @pytest.mark.parametrize('op', [operator.add, operator.sub,
+                                    operator.mul, operator.floordiv,
+                                    operator.truediv, operator.pow])
+    def test_arithmetic_with_frame_or_series(self, op):
+        # check that we return NotImplemented when operating with Series
+        # or DataFrame
+        index = pd.RangeIndex(5)
+        other = pd.Series(np.random.randn(5))
+
+        expected = op(pd.Series(index), other)
+        result = op(index, other)
+        tm.assert_series_equal(result, expected)
+
+        other = pd.DataFrame(np.random.randn(2, 5))
+        expected = op(pd.DataFrame([index, index]), other)
+        result = op(index, other)
+        tm.assert_frame_equal(result, expected)
+
+    def test_numeric_compat2(self):
+        # validate that we are handling the RangeIndex overrides to numeric ops
+        # and returning RangeIndex where possible
+
+        idx = pd.RangeIndex(0, 10, 2)
+
+        result = idx * 2
+        expected = pd.RangeIndex(0, 20, 4)
+        tm.assert_index_equal(result, expected, exact=True)
+
+        result = idx + 2
+        expected = pd.RangeIndex(2, 12, 2)
+        tm.assert_index_equal(result, expected, exact=True)
+
+        result = idx - 2
+        expected = pd.RangeIndex(-2, 8, 2)
+        tm.assert_index_equal(result, expected, exact=True)
+
+        # truediv under PY3
+        result = idx / 2
+
+        if PY3:
+            expected = pd.RangeIndex(0, 5, 1).astype('float64')
+        else:
+            expected = pd.RangeIndex(0, 5, 1)
+        tm.assert_index_equal(result, expected, exact=True)
+
+        result = idx / 4
+        expected = pd.RangeIndex(0, 10, 2) / 4
+        tm.assert_index_equal(result, expected, exact=True)
+
+        result = idx // 1
+        expected = idx
+        tm.assert_index_equal(result, expected, exact=True)
+
+        # __mul__
+        result = idx * idx
+        expected = Index(idx.values * idx.values)
+        tm.assert_index_equal(result, expected, exact=True)
+
+        # __pow__
+        idx = pd.RangeIndex(0, 1000, 2)
+        result = idx ** 2
+        expected = idx._int64index ** 2
+        tm.assert_index_equal(Index(result.values), expected, exact=True)
+
+        # __floordiv__
+        cases_exact = [
+            (pd.RangeIndex(0, 1000, 2), 2, pd.RangeIndex(0, 500, 1)),
+            (pd.RangeIndex(-99, -201, -3), -3, pd.RangeIndex(33, 67, 1)),
+            (pd.RangeIndex(0, 1000, 1), 2,
+             pd.RangeIndex(0, 1000, 1)._int64index // 2),
+            (pd.RangeIndex(0, 100, 1), 2.0,
+             pd.RangeIndex(0, 100, 1)._int64index // 2.0),
+            (pd.RangeIndex(0), 50, pd.RangeIndex(0)),
+            (pd.RangeIndex(2, 4, 2), 3, pd.RangeIndex(0, 1, 1)),
+            (pd.RangeIndex(-5, -10, -6), 4, pd.RangeIndex(-2, -1, 1)),
+            (pd.RangeIndex(-100, -200, 3), 2, pd.RangeIndex(0))]
+        for idx, div, expected in cases_exact:
+            tm.assert_index_equal(idx // div, expected, exact=True)

pandas/tests/arithmetic/test_object.py

@@ -73,6 +73,21 @@ def test_more_na_comparisons(self, dtype):
 
 class TestArithmetic(object):
 
+    # TODO: parametrize
+    def test_pow_ops_object(self):
+        # GH#22922
+        # pow is weird with masking & 1, so testing here
+        a = Series([1, np.nan, 1, np.nan], dtype=object)
+        b = Series([1, np.nan, np.nan, 1], dtype=object)
+        result = a ** b
+        expected = Series(a.values ** b.values, dtype=object)
+        tm.assert_series_equal(result, expected)
+
+        result = b ** a
+        expected = Series(b.values ** a.values, dtype=object)
+
+        tm.assert_series_equal(result, expected)
+
     @pytest.mark.parametrize("op", [operator.add, ops.radd])
     @pytest.mark.parametrize("other", ["category", "Int64"])
     def test_add_extension_scalar(self, other, box, op):