Merge pull request #157 from scipy/printed_arrays

ENH: more robust detection of printed arrays

Author: ev-br

scipy_doctest/impl.py

@@ -220,6 +220,30 @@ def try_convert_namedtuple(got):
     return got_again
 
 
+def try_convert_printed_array(got):
+    """Printed arrays: reinsert commas.
+    """
+    # a minimal version is `s_got = ", ".join(got[1:-1].split())`
+    # but it fails if there's a space after the opening bracket: "[ 0 1 2 ]"
+    # For 2D arrays, split into rows, drop spurious entries, then reassemble.
+    if not got.startswith('['):
+        return got
+
+    g1 = got[1:-1]  # strip outer "[...]"-s
+    rows = [x for x in g1.split("[") if x]
+    rows2 = [", ".join(row.split()) for row in rows]
+
+    if got.startswith("[["):
+        # was a 2D array, restore the opening brackets in rows; XXX clean up
+        rows3 = ["[" + row for row in rows2]
+    else:
+        rows3 = rows2
+
+    # add back the outer brackets
+    s_got = "[" + ", ".join(rows3) + "]"
+    return s_got
+
+
 def has_masked(got):
     return 'masked_array' in got and '--' in got
 
@@ -281,8 +305,9 @@ def check_output(self, want, got, optionflags):
             cond = (s_want.startswith("[") and s_want.endswith("]") and
                     s_got.startswith("[") and s_got.endswith("]"))
             if cond:
-                s_want = ", ".join(s_want[1:-1].split())
-                s_got = ", ".join(s_got[1:-1].split())
+                s_want = try_convert_printed_array(s_want)
+                s_got = try_convert_printed_array(s_got)
+
                 return self.check_output(s_want, s_got, optionflags)
 
             #handle array abbreviation for n-dimensional arrays, n >= 1

scipy_doctest/tests/module_cases.py

@@ -44,6 +44,32 @@ def func3():
     """
 
 
+def func_printed_arrays():
+    """
+    Check various ways handling of printed arrays can go wrong.
+
+    >>> import numpy as np
+    >>> a = np.arange(8).reshape(2, 4) / 3
+    >>> print(a)  # numpy 1.26.4
+    [[0.         0.33333333 0.66666667 1.        ]
+     [1.33333333 1.66666667 2.         2.33333333]]
+
+    >>> print(a)   # add spaces (older repr?)
+    [[ 0.         0.33333333 0.66666667 1.         ]
+     [ 1.33333333 1.66666667 2.         2.33333333 ]]
+
+    Also check 1D arrays
+    >>> a1 = np.arange(3)
+    >>> print(a1)
+    [0 1 2]
+    >>> print(a1)
+    [ 0 1 2]
+    >>> print(a1)
+    [ 0 1 2 ]
+
+    """
+
+
 def func4():
     """
     Test `# may vary` markers : these should not break doctests (but the code