Repository: quantopian/zipline
File: zipline/testing/predicates.py (Python)

Code before:
def assert_set_equal(result, expected, path=(), msg='', **kwargs):
    _check_sets(
        result,
        expected,
        msg,
        path,
        'element',
    )


@assert_equal.register(np.ndarray, np.ndarray)
def assert_array_equal(result,
                       expected,
                       path=(),
                       msg='',
                       array_verbose=True,
                       array_decimal=None,
                       **kwargs):
    f = (
        np.testing.assert_array_equal
        if array_decimal is None else
        partial(np.testing.assert_array_almost_equal, decimal=array_decimal)
    )

    is_datetime = False
    if result.dtype.kind == 'M' and expected.dtype.kind == 'M':
        is_datetime = True
        result = result.astype('int64')
        expected = expected.astype('int64')

    try:
        f(
            result,
            expected,
            verbose=array_verbose,
            err_msg=msg,
        )
    except AssertionError as e:
        if is_datetime:
            raise AssertionError(
                "\nArrays are not equal\n\n"
                "result: %s\n"
                "expected: %s\n"
                "%s" % (
                    np.array_repr(result.astype('M8[ns]')),
                    np.array_repr(expected.astype('M8[ns]')),
                    _fmt_path(path),
                )

            )
        raise AssertionError('\n'.join((str(e), _fmt_path(path))))

Review comment (suggestion):
You probably also want to handle `'m'` (timedelta) the same way. I'd write that check as:

```
if result.dtype.kind in 'mM' and expected.dtype == result.dtype:
```

Code after:
@assert_equal.register(set, set)
def assert_set_equal(result, expected, path=(), msg='', **kwargs):
    _check_sets(
        result,
        expected,
        msg,
        path,
        'element',
    )


@assert_equal.register(np.ndarray, np.ndarray)
def assert_array_equal(result,
                       expected,
                       path=(),
                       msg='',
                       array_verbose=True,
                       array_decimal=None,
                       **kwargs):
    result_dtype = result.dtype
    expected_dtype = expected.dtype

    if result_dtype.kind in 'mM' and expected_dtype.kind in 'mM':
        assert result_dtype == expected_dtype, (
            "\nType mismatch:\n\n"
            "result dtype: %s\n"
            "expected dtype: %s\n%s"
            % (result_dtype, expected_dtype, _fmt_path(path))
        )
        f = partial(
            np.testing.utils.assert_array_compare,
            compare_datetime_arrays,
            header='Arrays are not equal',
        )
    elif array_decimal is not None:
        f = partial(
            np.testing.assert_array_almost_equal,
            decimal=array_decimal,
        )
    else:
        f = np.testing.assert_array_equal

    try:
        f(
            result,
            expected,
            verbose=array_verbose,
            err_msg=msg,
        )
    except AssertionError as e:
        raise AssertionError('\n'.join((str(e), _fmt_path(path))))