Merge pull request #15997 from panpiort8/array_pretty_print

ENH: improve printing of arrays with multi-line reprs

2 files changed, 106 insertions, 4 deletions

diff --git a/numpy/core/arrayprint.py b/numpy/core/arrayprint.py
index 456ef76f0..5d9642ea8 100644
--- a/numpy/core/arrayprint.py
+++ b/numpy/core/arrayprint.py
@@ -695,7 +695,7 @@ def array2string(a, max_line_width=None, precision=None,
 def _extendLine(s, line, word, line_width, next_line_prefix, legacy):
     needs_wrap = len(line) + len(word) > line_width
     if legacy != '1.13':
-        s# don't wrap lines if it won't help
+        # don't wrap lines if it won't help
         if len(line) <= len(next_line_prefix):
             needs_wrap = False
 
@@ -706,6 +706,33 @@ def _extendLine(s, line, word, line_width, next_line_prefix, legacy):
     return s, line
 
 
+def _extendLine_pretty(s, line, word, line_width, next_line_prefix, legacy):
+    """
+    Extends line with nicely formatted (possibly multi-line) string ``word``.
+    """
+    words = word.splitlines()
+    if len(words) == 1 or legacy == '1.13':
+        return _extendLine(s, line, word, line_width, next_line_prefix, legacy)
+
+    max_word_length = max(len(word) for word in words)
+    if (len(line) + max_word_length > line_width and
+            len(line) > len(next_line_prefix)):
+        s += line.rstrip() + '\n'
+        line = next_line_prefix + words[0]
+        indent = next_line_prefix
+    else:
+        indent = len(line)*' '
+        line += words[0]
+
+    for word in words[1::]:
+        s += line.rstrip() + '\n'
+        line = indent + word
+
+    suffix_length = max_word_length - len(words[-1])
+    line += suffix_length*' '
+
+    return s, line
+
 def _formatArray(a, format_function, line_width, next_line_prefix,
                  separator, edge_items, summary_insert, legacy):
     """formatArray is designed for two modes of operation:
@@ -758,7 +785,7 @@ def _formatArray(a, format_function, line_width, next_line_prefix,
             line = hanging_indent
             for i in range(leading_items):
                 word = recurser(index + (i,), next_hanging_indent, next_width)
-                s, line = _extendLine(
+                s, line = _extendLine_pretty(
                     s, line, word, elem_width, hanging_indent, legacy)
                 line += separator
 
@@ -772,7 +799,7 @@ def _formatArray(a, format_function, line_width, next_line_prefix,
 
             for i in range(trailing_items, 1, -1):
                 word = recurser(index + (-i,), next_hanging_indent, next_width)
-                s, line = _extendLine(
+                s, line = _extendLine_pretty(
                     s, line, word, elem_width, hanging_indent, legacy)
                 line += separator
 
@@ -780,7 +807,7 @@ def _formatArray(a, format_function, line_width, next_line_prefix,
                 # width of the separator is not considered on 1.13
                 elem_width = curr_width
             word = recurser(index + (-1,), next_hanging_indent, next_width)
-            s, line = _extendLine(
+            s, line = _extendLine_pretty(
                 s, line, word, elem_width, hanging_indent, legacy)
 
             s += line
diff --git a/numpy/core/tests/test_arrayprint.py b/numpy/core/tests/test_arrayprint.py
index e29217461..a2703d81b 100644
--- a/numpy/core/tests/test_arrayprint.py
+++ b/numpy/core/tests/test_arrayprint.py
@@ -395,6 +395,81 @@ class TestArray2String:
             "[ 'xxxxx']"
         )
 
+    def test_multiline_repr(self):
+        class MultiLine:
+            def __repr__(self):
+                return "Line 1\nLine 2"
+
+        a = np.array([[None, MultiLine()], [MultiLine(), None]])
+
+        assert_equal(
+            np.array2string(a),
+            '[[None Line 1\n'
+            '       Line 2]\n'
+            ' [Line 1\n'
+            '  Line 2 None]]'
+        )
+        assert_equal(
+            np.array2string(a, max_line_width=5),
+            '[[None\n'
+            '  Line 1\n'
+            '  Line 2]\n'
+            ' [Line 1\n'
+            '  Line 2\n'
+            '  None]]'
+        )
+        assert_equal(
+            repr(a),
+            'array([[None, Line 1\n'
+            '              Line 2],\n'
+            '       [Line 1\n'
+            '        Line 2, None]], dtype=object)'
+        )
+
+        class MultiLineLong:
+            def __repr__(self):
+                return "Line 1\nLooooooooooongestLine2\nLongerLine 3"
+
+        a = np.array([[None, MultiLineLong()], [MultiLineLong(), None]])
+        assert_equal(
+            repr(a),
+            'array([[None, Line 1\n'
+            '              LooooooooooongestLine2\n'
+            '              LongerLine 3          ],\n'
+            '       [Line 1\n'
+            '        LooooooooooongestLine2\n'
+            '        LongerLine 3          , None]], dtype=object)'
+        )
+        assert_equal(
+            np.array_repr(a, 20),
+            'array([[None,\n'
+            '        Line 1\n'
+            '        LooooooooooongestLine2\n'
+            '        LongerLine 3          ],\n'
+            '       [Line 1\n'
+            '        LooooooooooongestLine2\n'
+            '        LongerLine 3          ,\n'
+            '        None]],\n'
+            '      dtype=object)'
+        )
+
+    def test_nested_array_repr(self):
+        a = np.empty((2, 2), dtype=object)
+        a[0, 0] = np.eye(2)
+        a[0, 1] = np.eye(3)
+        a[1, 0] = None
+        a[1, 1] = np.ones((3, 1))
+        assert_equal(
+            repr(a),
+            'array([[array([[1., 0.],\n'
+            '               [0., 1.]]), array([[1., 0., 0.],\n'
+            '                                  [0., 1., 0.],\n'
+            '                                  [0., 0., 1.]])],\n'
+            '       [None, array([[1.],\n'
+            '                     [1.],\n'
+            '                     [1.]])]], dtype=object)'
+        )
+
     @given(hynp.from_dtype(np.dtype("U")))
     def test_any_text(self, text):
         # This test checks that, given any value that can be represented in an