diff --git a/pandas/tests/arithmetic/conftest.py b/pandas/tests/arithmetic/conftest.py index 9ee5e05638978..e5d9c85691b7a 100644 --- a/pandas/tests/arithmetic/conftest.py +++ b/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 diff --git a/pandas/util/testing.py b/pandas/util/testing.py index 1fa77f5321038..1dc461e010e37 100644 --- a/pandas/util/testing.py +++ b/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