Improve prettyprinting of multidimensional numpy arrays.

prettyprinter/extras/numpy.py

@@ -1,6 +1,7 @@
 import ast
 from distutils.version import LooseVersion
 
+from ..doctypes import Concat, HARDLINE
 from ..prettyprinter import (
     register_pretty,
     pretty_call_alt,
@@ -10,6 +11,11 @@
 )
 
 
+class _ArrayWrapper:
+    def __init__(self, array):
+        self.array = array
+
+
 def pretty_ndarray(value, ctx):
     import numpy as np
     # numpy 1.14 added dtype_is_implied.
@@ -19,7 +25,7 @@ def pretty_ndarray(value, ctx):
         # Masked arrays, in particular, require their own logic.
         return repr(value)
     from numpy.core import arrayprint
-    args = (value.tolist(),)
+    args = (_ArrayWrapper(value),)
     kwargs = []
     dtype = value.dtype
     # This logic is extracted from arrayprint._array_repr_implementation.
@@ -31,6 +37,21 @@ def pretty_ndarray(value, ctx):
     return pretty_call_alt(ctx, type(value), args, kwargs)
 
 
+def pretty_arraywrapper(value, ctx):
+    import numpy as np
+    # array2string correctly aligns the items of the array, without adding
+    # "array(" in the front or the dtype info (which we handle ourselves) in
+    # the back.
+    s = np.array2string(value.array, separator=", ")
+    lines = s.split("\n")
+    if len(lines) == 1:
+        return s
+    else:
+        interspersed = [HARDLINE] * 2 * len(lines)
+        interspersed[1::2] = lines
+        return Concat(interspersed)
+
+
 def install():
     register_pretty("numpy.bool_")(pretty_bool)
 
@@ -46,3 +67,5 @@ def install():
         register_pretty("numpy." + name)(pretty_float)
 
     register_pretty("numpy.ndarray")(pretty_ndarray)
+    register_pretty("prettyprinter.extras.numpy._ArrayWrapper")(
+        pretty_arraywrapper)

tests/test_numpy.py

@@ -32,15 +32,22 @@ def test_numpy_bool_type():
     assert pformat(np.bool_(True)) == "numpy.bool_(True)"
 
 
-def test_array():
+def test_array_dtype():
     assert pformat(np.array([0, 1])) == "numpy.ndarray([0, 1])"
-    assert pformat(np.array([0., 1.])) == "numpy.ndarray([0.0, 1.0])"
+    assert pformat(np.array([0., 1.])) == "numpy.ndarray([0., 1.])"
     assert pformat(np.array([0, 1], dtype=np.uint8)) == "numpy.ndarray([0, 1], dtype='uint8')"
     assert pformat(np.array(["a", "b"])) == "numpy.ndarray(['a', 'b'], dtype='<U1')"
     assert pformat(np.array([("a", 1)], [("field1", str), ("field2", int)])) \
         == "numpy.ndarray([('', 1)], dtype=[('field1', '<U'), ('field2', '<i8')])"
 
 
+def test_array_2d():
+    assert pformat(np.arange(4).reshape((2, 2))) == \
+        "numpy.ndarray(\n    [[0, 1],\n     [2, 3]])"
+    assert pformat(np.arange(4, dtype=np.uint8).reshape((2, 2))) == \
+        "numpy.ndarray(\n    [[0, 1],\n     [2, 3]], dtype='uint8')"
+
+
 def test_masked_array():
     # Check that masked arrays don't go through the normal array pprinter
     # (they require their own algoithm, which is not implemented yet).