Reflect dimension in approx repr for numpy arrays.

Author: kalekundert

src/_pytest/python_api.py

@@ -82,14 +82,18 @@ class ApproxNumpy(ApproxBase):
     """
 
     def __repr__(self):
-        # It might be nice to rewrite this function to account for the
-        # shape of the array...
         import numpy as np
 
-        list_scalars = []
-        for x in np.ndindex(self.expected.shape):
-            list_scalars.append(self._approx_scalar(np.asscalar(self.expected[x])))
+        def recursive_map(f, x):
+            if isinstance(x, list):
+                return list(recursive_map(f, xi) for xi in x)
+            else:
+                return f(x)
 
+        list_scalars = recursive_map(
+                self._approx_scalar,
+                self.expected.tolist())
+                
         return "approx({!r})".format(list_scalars)
 
     if sys.version_info[0] == 2:

testing/python/approx.py

@@ -59,17 +59,19 @@ def test_repr_string(self, plus_minus):
             ),
         )
 
-    def test_repr_0d_array(self, plus_minus):
+    def test_repr_nd_array(self, plus_minus):
+        # Make sure that arrays of all different dimensions are repr'd 
+        # correctly.
         np = pytest.importorskip("numpy")
-        np_array = np.array(5.)
-        assert approx(np_array) == 5.0
-        string_expected = "approx([5.0 {} 5.0e-06])".format(plus_minus)
-
-        assert repr(approx(np_array)) == string_expected
-
-        np_array = np.array([5.])
-        assert approx(np_array) == 5.0
-        assert repr(approx(np_array)) == string_expected
+        examples = [
+                (np.array(5.), 'approx(5.0 {pm} 5.0e-06)'),
+                (np.array([5.]), 'approx([5.0 {pm} 5.0e-06])'),
+                (np.array([[5.]]), 'approx([[5.0 {pm} 5.0e-06]])'),
+                (np.array([[5., 6.]]), 'approx([[5.0 {pm} 5.0e-06, 6.0 {pm} 6.0e-06]])'),
+                (np.array([[5.], [6.]]), 'approx([[5.0 {pm} 5.0e-06], [6.0 {pm} 6.0e-06]])'),
+        ]
+        for np_array, repr_string in examples:
+            assert repr(approx(np_array)) == repr_string.format(pm=plus_minus)
 
     def test_operator_overloading(self):
         assert 1 == approx(1, rel=1e-6, abs=1e-12)