diff --git a/prettyprinter/extras/numpy.py b/prettyprinter/extras/numpy.py index 1768afe..e9aa1f6 100644 --- a/prettyprinter/extras/numpy.py +++ b/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) diff --git a/tests/test_numpy.py b/tests/test_numpy.py index 0c84b1f..4780050 100644 --- a/tests/test_numpy.py +++ b/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='