# -*- coding: utf-8 -*-

# -*- coding: utf-8 -*-

from __future__ import absolute_import

import dask.array.utils


def _assert_eq_nan(a, b, **kwargs):
    a = a.copy()
    b = b.copy()

    a_nan = (a != a)
    b_nan = (b != b)

    a[a_nan] = 0
    b[b_nan] = 0

    dask.array.utils.assert_eq(a_nan, b_nan, **kwargs)
    dask.array.utils.assert_eq(a, b, **kwargs)

Add _assert_eq_nan to compare arrays that have NaN
Add _assert_eq_nan to compare arrays that have NaN

As comparisons with NaN are false even if both values are NaN, using the
`assert_eq` does not work correctly in this case. To fix it, we add this
shim function around `assert_eq`. First we verify that they have the
same NaN values using a duck type friendly strategy (comparing the
arrays to themselves) and then comparing the masks of each array to the
other. Second we zero out all the NaN values and compare the resulting
arrays, which should now work as normal. This works just as well whether
the arrays are from NumPy, Dask, or some other library given they
support these basic functions.
