TST: implement tm.to_array, box_with_timedelta, box4 (pandas-dev#23748)

Author: jbrockmendel

pandas/tests/arithmetic/conftest.py

@@ -5,9 +5,25 @@
 import pandas as pd
 
 from pandas.compat import long
-from pandas.core.arrays import PeriodArray, DatetimeArrayMixin as DatetimeArray
+from pandas.core.arrays import (
+    PeriodArray, DatetimeArrayMixin as DatetimeArray,
+    TimedeltaArrayMixin as TimedeltaArray)
+import pandas.util.testing as tm
 
 
+# ------------------------------------------------------------------
+# Helper Functions
+
+def id_func(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x[0].__name__ + '-' + str(x[1])
+    else:
+        return x.__name__
+
+
+# ------------------------------------------------------------------
+
 @pytest.fixture(params=[1, np.array(1, dtype=np.int64)])
 def one(request):
     # zero-dim integer array behaves like an integer
@@ -137,7 +153,7 @@ def mismatched_freq(request):
 # ------------------------------------------------------------------
 
 @pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame],
-                ids=lambda x: x.__name__)
+                ids=id_func)
 def box(request):
     """
     Several array-like containers that should have effectively identical
@@ -150,7 +166,7 @@ def box(request):
                         pd.Series,
                         pytest.param(pd.DataFrame,
                                      marks=pytest.mark.xfail(strict=True))],
-                ids=lambda x: x.__name__)
+                ids=id_func)
 def box_df_fail(request):
     """
     Fixture equivalent to `box` fixture but xfailing the DataFrame case.
@@ -163,7 +179,7 @@ def box_df_fail(request):
                         (pd.DataFrame, False),
                         pytest.param((pd.DataFrame, True),
                                      marks=pytest.mark.xfail(strict=True))],
-                ids=lambda x: x[0].__name__ + '-' + str(x[1]))
+                ids=id_func)
 def box_transpose_fail(request):
     """
     Fixture similar to `box` but testing both transpose cases for DataFrame,
@@ -174,7 +190,7 @@ def box_transpose_fail(request):
 
 
 @pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, PeriodArray],
-                ids=lambda x: x.__name__)
+                ids=id_func)
 def box_with_period(request):
     """
     Like `box`, but specific to PeriodDtype for also testing PeriodArray
@@ -183,9 +199,28 @@ def box_with_period(request):
 
 
 @pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, DatetimeArray],
-                ids=lambda x: x.__name__)
+                ids=id_func)
 def box_with_datetime(request):
     """
     Like `box`, but specific to datetime64 for also testing DatetimeArray
     """
     return request.param
+
+
+@pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, TimedeltaArray],
+                ids=id_func)
+def box_with_timedelta(request):
+    """
+    Like `box`, but specific to timedelta64 for also testing TimedeltaArray
+    """
+    return request.param
+
+
+@pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, tm.to_array],
+                ids=id_func)
+def box_with_array(request):
+    """
+    Fixture to test behavior for Index, Series, DataFrame, and pandas Array
+    classes
+    """
+    return request.param

pandas/util/testing.py

@@ -24,9 +24,10 @@
     map, raise_with_traceback, range, string_types, u, unichr, zip)
 
 from pandas.core.dtypes.common import (
-    is_bool, is_categorical_dtype, is_datetimelike_v_numeric,
-    is_datetimelike_v_object, is_extension_array_dtype, is_interval_dtype,
-    is_list_like, is_number, is_sequence, needs_i8_conversion)
+    is_bool, is_categorical_dtype, is_datetime64_dtype, is_datetime64tz_dtype,
+    is_datetimelike_v_numeric, is_datetimelike_v_object,
+    is_extension_array_dtype, is_interval_dtype, is_list_like, is_number,
+    is_period_dtype, is_sequence, is_timedelta64_dtype, needs_i8_conversion)
 from pandas.core.dtypes.missing import array_equivalent
 
 import pandas as pd
@@ -37,7 +38,7 @@
 from pandas.core.algorithms import take_1d
 from pandas.core.arrays import (
     DatetimeArrayMixin as DatetimeArray, ExtensionArray, IntervalArray,
-    PeriodArray, period_array)
+    PeriodArray, TimedeltaArrayMixin as TimedeltaArray, period_array)
 import pandas.core.common as com
 
 from pandas.io.common import urlopen
@@ -1080,6 +1081,13 @@ def assert_datetime_array_equal(left, right, obj='DatetimeArray'):
     assert_attr_equal('tz', left, right, obj=obj)
 
 
+def assert_timedelta_array_equal(left, right, obj='TimedeltaArray'):
+    _check_isinstance(left, right, TimedeltaArray)
+    assert_numpy_array_equal(left._data, right._data,
+                             obj='{obj}._data'.format(obj=obj))
+    assert_attr_equal('freq', left, right, obj=obj)
+
+
 def raise_assert_detail(obj, message, left, right, diff=None):
     __tracebackhide__ = True
 
@@ -1579,6 +1587,8 @@ def assert_equal(left, right, **kwargs):
         assert_period_array_equal(left, right, **kwargs)
     elif isinstance(left, DatetimeArray):
         assert_datetime_array_equal(left, right, **kwargs)
+    elif isinstance(left, TimedeltaArray):
+        assert_timedelta_array_equal(left, right, **kwargs)
     elif isinstance(left, ExtensionArray):
         assert_extension_array_equal(left, right, **kwargs)
     elif isinstance(left, np.ndarray):
@@ -1616,13 +1626,29 @@ def box_expected(expected, box_cls, transpose=True):
         expected = period_array(expected)
     elif box_cls is DatetimeArray:
         expected = DatetimeArray(expected)
+    elif box_cls is TimedeltaArray:
+        expected = TimedeltaArray(expected)
     elif box_cls is np.ndarray:
         expected = np.array(expected)
+    elif box_cls is to_array:
+        expected = to_array(expected)
     else:
         raise NotImplementedError(box_cls)
     return expected
 
 
+def to_array(obj):
+    # temporary implementation until we get pd.array in place
+    if is_period_dtype(obj):
+        return period_array(obj)
+    elif is_datetime64_dtype(obj) or is_datetime64tz_dtype(obj):
+        return DatetimeArray(obj)
+    elif is_timedelta64_dtype(obj):
+        return TimedeltaArray(obj)
+    else:
+        return np.array(obj)
+
+
 # -----------------------------------------------------------------------------
 # Sparse