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

Code before:
                       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))))


@assert_equal.register(LabelArray, LabelArray)
def assert_labelarray_equal(result, expected, path=(), **kwargs):
    assert_equal(
        result.categories,
        expected.categories,
        path=path + ('.categories',),
        **kwargs
    )
    assert_equal(
        result.as_int_array(),
        expected.as_int_array(),

Review comment (suggestion):
Hmm. this feels kinda gross. I feel like it might be cleaner to just delegate to an `assert_datetime_array_equal` method so that we don't have to keep the flag around.

Code after:
                       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))))


@assert_equal.register(LabelArray, LabelArray)
def assert_labelarray_equal(result, expected, path=(), **kwargs):
    assert_equal(
        result.categories,
        expected.categories,
        path=path + ('.categories',),
        **kwargs
    )
    assert_equal(
        result.as_int_array(),
        expected.as_int_array(),