TST: Test cleanup, parametrization for datetime64 arithmetic tests (#23681)

Author: jbrockmendel

pandas/conftest.py

@@ -1,10 +1,12 @@
 import importlib
 import os
 
+from dateutil.tz import tzutc
 import hypothesis
 from hypothesis import strategies as st
 import numpy as np
 import pytest
+from pytz import utc
 
 from pandas.compat import PY3
 import pandas.util._test_decorators as td
@@ -243,6 +245,20 @@ def datetime_tz_utc():
     return timezone.utc
 
 
+utc_objs = ['utc', utc, tzutc()]
+if PY3:
+    from datetime import timezone
+    utc_objs.append(timezone.utc)
+
+
+@pytest.fixture(params=utc_objs)
+def utc_fixture(request):
+    """
+    Fixture to provide variants of UTC timezone strings and tzinfo objects
+    """
+    return request.param
+
+
 @pytest.fixture(params=['inner', 'outer', 'left', 'right'])
 def join_type(request):
     """

pandas/core/arrays/datetimes.py

@@ -110,7 +110,7 @@ def wrapper(self, other):
                 # string that cannot be parsed to Timestamp
                 return ops.invalid_comparison(self, other, op)
 
-            result = meth(self, other)
+            result = op(self.asi8, other.view('i8'))
             if isna(other):
                 result.fill(nat_result)
         elif lib.is_scalar(other):
@@ -208,9 +208,6 @@ def _simple_new(cls, values, freq=None, tz=None, **kwargs):
         return result
 
     def __new__(cls, values, freq=None, tz=None, dtype=None):
-        if tz is None and hasattr(values, 'tz'):
-            # e.g. DatetimeIndex
-            tz = values.tz
 
         if freq is None and hasattr(values, "freq"):
             # i.e. DatetimeArray, DatetimeIndex
@@ -221,9 +218,16 @@ def __new__(cls, values, freq=None, tz=None, dtype=None):
         # if dtype has an embedded tz, capture it
         tz = dtl.validate_tz_from_dtype(dtype, tz)
 
+        if isinstance(values, ABCSeries):
+            # extract to ndarray or DatetimeIndex
+            values = values._values
+
         if isinstance(values, DatetimeArrayMixin):
             # extract nanosecond unix timestamps
+            if tz is None:
+                tz = values.tz
             values = values.asi8
+
         if values.dtype == 'i8':
             values = values.view('M8[ns]')
 

pandas/tests/arithmetic/test_datetime64.py

@@ -25,6 +25,7 @@
 from pandas import (
     Timestamp, Timedelta, Period, Series, date_range, NaT,
     DatetimeIndex, TimedeltaIndex)
+from pandas.core.arrays import DatetimeArrayMixin as DatetimeArray
 
 
 # ------------------------------------------------------------------
@@ -158,12 +159,16 @@ def test_dt64_ser_cmp_date_warning(self):
         assert "a TypeError will be raised" in str(m[0].message)
 
     @pytest.mark.skip(reason="GH#21359")
-    def test_dt64ser_cmp_date_invalid(self):
+    def test_dt64ser_cmp_date_invalid(self, box_with_datetime):
         # GH#19800 datetime.date comparison raises to
         # match DatetimeIndex/Timestamp.  This also matches the behavior
         # of stdlib datetime.datetime
-        ser = pd.Series(pd.date_range('20010101', periods=10), name='dates')
+        box = box_with_datetime
+
+        ser = pd.date_range('20010101', periods=10)
         date = ser.iloc[0].to_pydatetime().date()
+
+        ser = tm.box_expected(ser, box)
         assert not (ser == date).any()
         assert (ser != date).all()
         with pytest.raises(TypeError):
@@ -225,22 +230,37 @@ def test_timestamp_compare_series(self, left, right):
         result = right_f(pd.Timestamp("nat"), s_nat)
         tm.assert_series_equal(result, expected)
 
-    def test_timestamp_equality(self):
+    def test_dt64arr_timestamp_equality(self, box_with_datetime):
         # GH#11034
+        box = box_with_datetime
+        xbox = box if box not in [pd.Index, DatetimeArray] else np.ndarray
+
         ser = pd.Series([pd.Timestamp('2000-01-29 01:59:00'), 'NaT'])
+        ser = tm.box_expected(ser, box)
+
         result = ser != ser
-        tm.assert_series_equal(result, pd.Series([False, True]))
+        expected = tm.box_expected([False, True], xbox)
+        tm.assert_equal(result, expected)
+
         result = ser != ser[0]
-        tm.assert_series_equal(result, pd.Series([False, True]))
+        expected = tm.box_expected([False, True], xbox)
+        tm.assert_equal(result, expected)
+
         result = ser != ser[1]
-        tm.assert_series_equal(result, pd.Series([True, True]))
+        expected = tm.box_expected([True, True], xbox)
+        tm.assert_equal(result, expected)
 
         result = ser == ser
-        tm.assert_series_equal(result, pd.Series([True, False]))
+        expected = tm.box_expected([True, False], xbox)
+        tm.assert_equal(result, expected)
+
         result = ser == ser[0]
-        tm.assert_series_equal(result, pd.Series([True, False]))
+        expected = tm.box_expected([True, False], xbox)
+        tm.assert_equal(result, expected)
+
         result = ser == ser[1]
-        tm.assert_series_equal(result, pd.Series([False, False]))
+        expected = tm.box_expected([False, False], xbox)
+        tm.assert_equal(result, expected)
 
 
 class TestDatetimeIndexComparisons(object):
@@ -629,7 +649,7 @@ def test_dti_cmp_object_dtype(self):
 # Arithmetic
 
 class TestFrameArithmetic(object):
-    def test_dt64arr_sub_dtscalar(self, box):
+    def test_dt64arr_sub_timestamp(self, box):
         # GH#8554, GH#22163 DataFrame op should _not_ return dt64 dtype
         idx = pd.date_range('2013-01-01', periods=3)
         idx = tm.box_expected(idx, box)
@@ -643,28 +663,39 @@ def test_dt64arr_sub_dtscalar(self, box):
         result = idx - ts
         tm.assert_equal(result, expected)
 
-    def test_df_sub_datetime64_not_ns(self):
+    def test_dt64arr_sub_datetime64_not_ns(self, box):
         # GH#7996, GH#22163 ensure non-nano datetime64 is converted to nano
-        df = pd.DataFrame(pd.date_range('20130101', periods=3))
+        #  for DataFrame operation
+
+        dti = pd.date_range('20130101', periods=3)
+        dtarr = tm.box_expected(dti, box)
+
         dt64 = np.datetime64('2013-01-01')
         assert dt64.dtype == 'datetime64[D]'
-        res = df - dt64
-        expected = pd.DataFrame([pd.Timedelta(days=0), pd.Timedelta(days=1),
-                                 pd.Timedelta(days=2)])
-        tm.assert_frame_equal(res, expected)
+
+        expected = pd.TimedeltaIndex(['0 Days', '1 Day', '2 Days'])
+        expected = tm.box_expected(expected, box)
+
+        result = dtarr - dt64
+        tm.assert_equal(result, expected)
 
 
 class TestTimestampSeriesArithmetic(object):
 
-    def test_timestamp_sub_series(self):
-        ser = pd.Series(pd.date_range('2014-03-17', periods=2, freq='D',
-                                      tz='US/Eastern'))
+    def test_dt64arr_sub_timestamp(self, box):
+        ser = pd.date_range('2014-03-17', periods=2, freq='D',
+                            tz='US/Eastern')
         ts = ser[0]
 
+        # FIXME: transpose raises ValueError
+        ser = tm.box_expected(ser, box, transpose=False)
+
         delta_series = pd.Series([np.timedelta64(0, 'D'),
                                   np.timedelta64(1, 'D')])
-        tm.assert_series_equal(ser - ts, delta_series)
-        tm.assert_series_equal(ts - ser, -delta_series)
+        expected = tm.box_expected(delta_series, box, transpose=False)
+
+        tm.assert_equal(ser - ts, expected)
+        tm.assert_equal(ts - ser, -expected)
 
     def test_dt64ser_sub_datetime_dtype(self):
         ts = Timestamp(datetime(1993, 1, 7, 13, 30, 00))
@@ -722,20 +753,23 @@ def check(get_ser, test_ser):
         if op_str not in ['__add__', '__radd__', '__sub__', '__rsub__']:
             check(dt2, td2)
 
-    @pytest.mark.parametrize('klass', [Series, pd.Index])
-    def test_sub_datetime64_not_ns(self, klass):
-        # GH#7996
+    def test_sub_datetime64_not_ns(self, box):
+        # GH#7996 operation with non-nano datetime64 scalar
         dt64 = np.datetime64('2013-01-01')
         assert dt64.dtype == 'datetime64[D]'
 
-        obj = klass(date_range('20130101', periods=3))
-        res = obj - dt64
-        expected = klass([Timedelta(days=0), Timedelta(days=1),
-                          Timedelta(days=2)])
-        tm.assert_equal(res, expected)
+        obj = date_range('20130101', periods=3)
+        obj = tm.box_expected(obj, box)
+
+        expected = TimedeltaIndex([Timedelta(days=0), Timedelta(days=1),
+                                   Timedelta(days=2)])
+        expected = tm.box_expected(expected, box)
+
+        result = obj - dt64
+        tm.assert_equal(result, expected)
 
-        res = dt64 - obj
-        tm.assert_equal(res, -expected)
+        result = dt64 - obj
+        tm.assert_equal(result, -expected)
 
     def test_sub_single_tz(self):
         # GH12290
@@ -1438,34 +1472,43 @@ def test_sub_dti_dti(self):
         result = dti2 - dti1
         tm.assert_index_equal(result, expected)
 
-    @pytest.mark.parametrize('freq', [None, 'D'])
-    def test_sub_period(self, freq, box_with_datetime):
+    @pytest.mark.parametrize('dti_freq', [None, 'D'])
+    def test_dt64arr_add_sub_period(self, dti_freq, box_with_datetime):
         # GH#13078
         # not supported, check TypeError
         p = pd.Period('2011-01-01', freq='D')
 
-        idx = pd.DatetimeIndex(['2011-01-01', '2011-01-02'], freq=freq)
+        idx = pd.DatetimeIndex(['2011-01-01', '2011-01-02'], freq=dti_freq)
         idx = tm.box_expected(idx, box_with_datetime)
 
+        with pytest.raises(TypeError):
+            idx + p
+        with pytest.raises(TypeError):
+            p + idx
         with pytest.raises(TypeError):
             idx - p
-
         with pytest.raises(TypeError):
             p - idx
 
-    @pytest.mark.parametrize('op', [operator.add, ops.radd,
-                                    operator.sub, ops.rsub])
     @pytest.mark.parametrize('pi_freq', ['D', 'W', 'Q', 'H'])
     @pytest.mark.parametrize('dti_freq', [None, 'D'])
-    def test_dti_sub_pi(self, dti_freq, pi_freq, op, box):
+    def test_dti_add_sub_pi(self, dti_freq, pi_freq,
+                            box_with_datetime, box_with_period):
         # GH#20049 subtracting PeriodIndex should raise TypeError
         dti = pd.DatetimeIndex(['2011-01-01', '2011-01-02'], freq=dti_freq)
         pi = dti.to_period(pi_freq)
 
-        dti = tm.box_expected(dti, box)
-        # TODO: Also box pi?
+        dtarr = tm.box_expected(dti, box_with_datetime)
+        parr = tm.box_expected(pi, box_with_period)
+
+        with pytest.raises(TypeError):
+            dtarr + parr
+        with pytest.raises(TypeError):
+            parr + dtarr
+        with pytest.raises(TypeError):
+            dtarr - parr
         with pytest.raises(TypeError):
-            op(dti, pi)
+            parr - dtarr
 
     # -------------------------------------------------------------------
     # TODO: Most of this block is moved from series or frame tests, needs