Better introspection for comparing numpy arrays

[timhoffm]

With plain python lists I can do:

a = [1, 2, 3]
b = [1, 2, 4]
assert a == b

That does not work for numpy arrays

        a = np.array([1, 2, 3])
        b = np.array([1, 2, 4])
>       assert a == b
E       ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

There are two ways out:

• Using numpy.testing

  np.testing.assert_equal(a, b)
  

  This has the advantage of giving good context information on the error.
  But it is far from the simple assert assert a == b paradigm in pytest.

• Alternatively we can aggregate the the check to a single bool

  assert np.all(a == b)
  

  This looks much nicer, but pytest cannot tell anymore where the error is:

  >       assert np.all(a == b)
  E       assert False
  E        +  where False = <function all at 0x7f57cd65eea0>(array([1, 2, 3]) == array([1, 2, 4])
  E        +    where <function all at 0x7f57cd65eea0> = np.all
  

Question: Would it in principle be possible with pytest's introspection capabilities to resolve the contents of np.all() and give a better error message? In a frist step maybe something like

assert np.all(a == b)
  At index 2: False

This of course, would require an assumption on the semantics of np.all.

Alternative idea: Could pytest be made to rewrite assert a == b to np.testing.assert_equal(a, b) if either a or b is a numpy array? Possibly not a good idea because it introduces a different semantic compared to a plain a == b.

[asottile]

== is not an equality check for numpy arrays

>>> x = np.array([1, 2, 3])
>>> y = np.array([1, 2, 3])
>>> x == y
array([ True,  True,  True])

[Zac-HD]

IMO the idiomatic solution is to use np.testing, and if you want more detailed feedback the right place to generate more detailed messages would be upstream in Numpy.

(we've run into similar issues in Hypothesis; numpy arrays are just complex enough to need special handling and papering over the differences only makes it worse in the long run 😭)

[twmr]

Does pytest allow to write e.g. a numpy testing plugin that extends the assert statement bytecode rewriter? This plugin would also need to provide a new version of the pytest.approx function, which exposes the functionality of the numpy.testing.assert_* functions.

[nicoddemus]

numpy testing plugin that extends the assert statement bytecode rewriter?

That's an interesting proposal. We would need to figure out an API and hooks in order to do that though. This is interesting as it would allow all kinds of extensions, like the one proposed here and #3479.

[asottile]

even with that, I don't think the one proposed here is a good proposal, numpy overrides == to implement a projection -- if that code were suddenly in a test and didn't act as a projection that would be pretty surprising (maybe not as surprising as numpy doing that in the first place, but still a deviation from how the code actually executes)

[timhoffm]

I agree that rewriting a == b is a bit dangerous. OTOH the specific expression assert a == b does currently not have a defined meaning. But nevertheless it's probably not good to make a == b context dependent. And it still does not solve the approx/allclose cases. This was just an idea i wanted to add to the discussion to have the space of possible solutions covered.

Independently of this is the question if pytest could introspect assert np.all(a == b) or assert np.allclose(a, b). Essentially, that could replace the need for explicit numpy.testing.* functions. Instead of

assert a.dtype == b.dtype
nptest.assert_allclose(a, b)

one could then write

assert a.dtype == b.dtype
assert np.allclose(a, b)

which restores the pattern of starting all test conditions with assert.

[samueljsb]

The error is a numpy error, which you'd get when asserting the truthiness of an array regardless of whether you're using pytest or not (this falls out of the numpy implementation of __eq__ for arrays). Is pytest core really the right place to change that behaviour? The error message suggests a solution: use assert all(a==b).

[timhoffm]

@samueljsb As I've written in the original message assert np.all(a == b) is one way of writing this. But it has the disadvantage of a not quite instructive error message.

Additional problem: If the dimensions are not the same, numpy currently just returns False for a == b. Thus, tests can fail with an obscure error message:

>       assert all(a == b)
E       TypeError: 'bool' object is not iterable

Fundamental conclusion: Numpy arrays are too complex for attempting any simple assert * with introspection from the pytest side. The comparison logic and message must come from numpy.

These are implemented in numpy.testing. At present

assert a.dtype == b.dtype
nptest.assert_allclose(a, b)

is the right way to go.

Possible enhancement: A plugin system for assert statement bytecode rewriting as proposed in #5347 (comment) could allow to write

assert a.dtype == b.dtype
assert np.allclose(a, b)

The rewrite rule assert np.allclose --> nptest.assert_allclose would have to be provided by numpy. This way users would not have to learn about np.testing explicitly. They would use normal numpy comparison functions but would still receive more instructive error messages.

[MachinesJesus]

Informative dev issue. Thanks!

I ended up just doing:

assert np.testing.assert_allclose(a, expected_a) is None

seems to work for now, don't know if you guys have fixed it or have a better patch. Will be using this a lot.