From c34dde7a3f582ddb57560ea08cfe00b84be31e16 Mon Sep 17 00:00:00 2001
From: Tadeu Manoel <e.tadeu@gmail.com>
Date: Wed, 14 Mar 2018 15:29:40 -0300
Subject: [PATCH] Add support for pytest.approx comparisons between array and
 scalar

---
 _pytest/python_api.py    | 25 +++++++++++++++++++++----
 changelog/3312.feature   |  1 +
 testing/python/approx.py | 11 +++++++++++
 3 files changed, 33 insertions(+), 4 deletions(-)
 create mode 100644 changelog/3312.feature

diff --git a/_pytest/python_api.py b/_pytest/python_api.py
index 3dce7f6b4..4b428322a 100644
--- a/_pytest/python_api.py
+++ b/_pytest/python_api.py
@@ -31,6 +31,10 @@ class ApproxBase(object):
     or sequences of numbers.
     """
 
+    # Tell numpy to use our `__eq__` operator instead of its when left side in a numpy array but right side is
+    # an instance of ApproxBase
+    __array_ufunc__ = None
+
     def __init__(self, expected, rel=None, abs=None, nan_ok=False):
         self.expected = expected
         self.abs = abs
@@ -89,7 +93,7 @@ class ApproxNumpy(ApproxBase):
         except:  # noqa
             raise TypeError("cannot compare '{0}' to numpy.ndarray".format(actual))
 
-        if actual.shape != self.expected.shape:
+        if not np.isscalar(self.expected) and actual.shape != self.expected.shape:
             return False
 
         return ApproxBase.__eq__(self, actual)
@@ -100,8 +104,13 @@ class ApproxNumpy(ApproxBase):
         # We can be sure that `actual` is a numpy array, because it's
         # casted in `__eq__` before being passed to `ApproxBase.__eq__`,
         # which is the only method that calls this one.
-        for i in np.ndindex(self.expected.shape):
-            yield actual[i], self.expected[i]
+
+        if np.isscalar(self.expected):
+            for i in np.ndindex(actual.shape):
+                yield actual[i], self.expected
+        else:
+            for i in np.ndindex(self.expected.shape):
+                yield actual[i], self.expected[i]
 
 
 class ApproxMapping(ApproxBase):
@@ -189,6 +198,8 @@ class ApproxScalar(ApproxBase):
         Return true if the given value is equal to the expected value within
         the pre-specified tolerance.
         """
+        if _is_numpy_array(actual):
+            return actual == ApproxNumpy(self.expected, self.abs, self.rel, self.nan_ok)
 
         # Short-circuit exact equality.
         if actual == self.expected:
@@ -308,12 +319,18 @@ def approx(expected, rel=None, abs=None, nan_ok=False):
         >>> {'a': 0.1 + 0.2, 'b': 0.2 + 0.4} == approx({'a': 0.3, 'b': 0.6})
         True
 
-    And ``numpy`` arrays::
+    ``numpy`` arrays::
 
         >>> import numpy as np                                                          # doctest: +SKIP
         >>> np.array([0.1, 0.2]) + np.array([0.2, 0.4]) == approx(np.array([0.3, 0.6])) # doctest: +SKIP
         True
 
+    And for a ``numpy`` array against a scalar::
+
+        >>> import numpy as np                                         # doctest: +SKIP
+        >>> np.array([0.1, 0.2]) + np.array([0.2, 0.1]) == approx(0.3) # doctest: +SKIP
+        True
+
     By default, ``approx`` considers numbers within a relative tolerance of
     ``1e-6`` (i.e. one part in a million) of its expected value to be equal.
     This treatment would lead to surprising results if the expected value was
diff --git a/changelog/3312.feature b/changelog/3312.feature
new file mode 100644
index 000000000..ffb4df8e9
--- /dev/null
+++ b/changelog/3312.feature
@@ -0,0 +1 @@
+``pytest.approx`` now accepts comparing a numpy array with a scalar.
diff --git a/testing/python/approx.py b/testing/python/approx.py
index 341e5fcff..b9d28aadb 100644
--- a/testing/python/approx.py
+++ b/testing/python/approx.py
@@ -391,3 +391,14 @@ class TestApprox(object):
         """
         with pytest.raises(TypeError):
             op(1, approx(1, rel=1e-6, abs=1e-12))
+
+    def test_numpy_array_with_scalar(self):
+        np = pytest.importorskip('numpy')
+
+        actual = np.array([1 + 1e-7, 1 - 1e-8])
+        expected = 1.0
+
+        assert actual == approx(expected, rel=5e-7, abs=0)
+        assert actual != approx(expected, rel=5e-8, abs=0)
+        assert approx(expected, rel=5e-7, abs=0) == actual
+        assert approx(expected, rel=5e-8, abs=0) != actual