Merge pull request #27482 from mhvk/arrayprint-shape-for-summarized

mattip · web-flow · commit fbffb8cc29eb · 2024-11-07T11:34:33.000+02:00
Show shape any time it cannot be inferred in repr
diff --git a/doc/release/upcoming_changes/27482.change.rst b/doc/release/upcoming_changes/27482.change.rst
@@ -0,0 +1,8 @@
+* The ``repr`` of arrays large enough to be summarized (i.e., where elements
+  are replaced with ``...``) now includes the ``shape`` of the array, similar
+  to what already was the case for arrays with zero size and non-obvious
+  shape. With this change, the shape is always given when it cannot be
+  inferred from the values.  Note that while written as ``shape=...``, this
+  argument cannot actually be passed in to the ``np.array`` constructor. If
+  you encounter problems, e.g., due to failing doctests, you can use the print
+  option ``legacy=2.1`` to get the old behaviour.
diff --git a/numpy/_core/arrayprint.py b/numpy/_core/arrayprint.py
@@ -83,12 +83,14 @@ def _make_options_dict(precision=None, threshold=None, edgeitems=None,
         options['legacy'] = 121
     elif legacy == '1.25':
         options['legacy'] = 125
+    elif legacy == '2.1':
+        options['legacy'] = 201
     elif legacy is None:
         pass  # OK, do nothing.
     else:
         warnings.warn(
             "legacy printing option can currently only be '1.13', '1.21', "
-            "'1.25', or `False`", stacklevel=3)
+            "'1.25', '2.1, or `False`", stacklevel=3)
 
     if threshold is not None:
         # forbid the bad threshold arg suggested by stack overflow, gh-12351
@@ -214,13 +216,16 @@ def set_printoptions(precision=None, threshold=None, edgeitems=None,
         that numeric scalars are printed without their type information, e.g.
         as ``3.0`` rather than ``np.float64(3.0)``.
 
+        If set to ``'2.1'``, shape information is not given when arrays are
+        summarized (i.e., multiple elements replaced with ``...``).
+
         If set to `False`, disables legacy mode.
 
         Unrecognized strings will be ignored with a warning for forward
         compatibility.
 
         .. versionchanged:: 1.22.0
-        .. versionchanged:: 2.0
+        .. versionchanged:: 2.2
 
     override_repr: callable, optional
         If set a passed function will be used for generating arrays' repr.
@@ -249,7 +254,7 @@ def set_printoptions(precision=None, threshold=None, edgeitems=None,
 
     >>> np.set_printoptions(threshold=5)
     >>> np.arange(10)
-    array([0, 1, 2, ..., 7, 8, 9])
+    array([0, 1, 2, ..., 7, 8, 9], shape=(10,))
 
     Small results can be suppressed:
 
@@ -282,7 +287,7 @@ def set_printoptions(precision=None, threshold=None, edgeitems=None,
 
     >>> with np.printoptions(precision=2, suppress=True, threshold=5):
     ...     np.linspace(0, 10, 10)
-    array([ 0.  ,  1.11,  2.22, ...,  7.78,  8.89, 10.  ])
+    array([ 0.  ,  1.11,  2.22, ...,  7.78,  8.89, 10.  ], shape=(10,))
 
     """
     _set_printoptions(precision, threshold, edgeitems, linewidth, suppress,
@@ -1578,39 +1583,41 @@ def _array_repr_implementation(
     else:
         class_name = "array"
 
-    skipdtype = dtype_is_implied(arr.dtype) and arr.size > 0
-
     prefix = class_name + "("
-    suffix = ")" if skipdtype else ","
-
     if (current_options['legacy'] <= 113 and
             arr.shape == () and not arr.dtype.names):
         lst = repr(arr.item())
-    elif arr.size > 0 or arr.shape == (0,):
+    else:
         lst = array2string(arr, max_line_width, precision, suppress_small,
-                           ', ', prefix, suffix=suffix)
-    else:  # show zero-length shape unless it is (0,)
-        lst = "[], shape=%s" % (repr(arr.shape),)
-
-    arr_str = prefix + lst + suffix
-
-    if skipdtype:
-        return arr_str
-
-    dtype_str = "dtype={})".format(dtype_short_repr(arr.dtype))
-
-    # compute whether we should put dtype on a new line: Do so if adding the
-    # dtype would extend the last line past max_line_width.
+                           ', ', prefix, suffix=")")
+
+    # Add dtype and shape information if these cannot be inferred from
+    # the array string.
+    extras = []
+    if (arr.size == 0 and arr.shape != (0,)
+            or current_options['legacy'] > 210
+            and arr.size > current_options['threshold']):
+        extras.append(f"shape={arr.shape}")
+    if not dtype_is_implied(arr.dtype) or arr.size == 0:
+        extras.append(f"dtype={dtype_short_repr(arr.dtype)}")
+
+    if not extras:
+        return prefix + lst + ")"
+
+    arr_str = prefix + lst + ","
+    extra_str = ", ".join(extras) + ")"
+    # compute whether we should put extras on a new line: Do so if adding the
+    # extras would extend the last line past max_line_width.
     # Note: This line gives the correct result even when rfind returns -1.
     last_line_len = len(arr_str) - (arr_str.rfind('\n') + 1)
     spacer = " "
     if current_options['legacy'] <= 113:
         if issubclass(arr.dtype.type, flexible):
-            spacer = '\n' + ' '*len(class_name + "(")
-    elif last_line_len + len(dtype_str) + 1 > max_line_width:
-        spacer = '\n' + ' '*len(class_name + "(")
+            spacer = '\n' + ' '*len(prefix)
+    elif last_line_len + len(extra_str) + 1 > max_line_width:
+        spacer = '\n' + ' '*len(prefix)
 
-    return arr_str + spacer + dtype_str
+    return arr_str + spacer + extra_str
 
 
 def _array_repr_dispatcher(
diff --git a/numpy/_core/tests/test_arrayprint.py b/numpy/_core/tests/test_arrayprint.py
@@ -346,7 +346,13 @@ def test_summarize_1d(self):
         assert_equal(str(A), strA)
 
         reprA = 'array([   0,    1,    2, ...,  998,  999, 1000])'
-        assert_equal(repr(A), reprA)
+        try:
+            np.set_printoptions(legacy='2.1')
+            assert_equal(repr(A), reprA)
+        finally:
+            np.set_printoptions(legacy=False)
+
+        assert_equal(repr(A), reprA.replace(')', ', shape=(1001,))'))
 
     def test_summarize_2d(self):
         A = np.arange(1002).reshape(2, 501)
@@ -356,6 +362,23 @@ def test_summarize_2d(self):
 
         reprA = 'array([[   0,    1,    2, ...,  498,  499,  500],\n' \
                 '       [ 501,  502,  503, ...,  999, 1000, 1001]])'
+        try:
+            np.set_printoptions(legacy='2.1')
+            assert_equal(repr(A), reprA)
+        finally:
+            np.set_printoptions(legacy=False)
+
+        assert_equal(repr(A), reprA.replace(')', ', shape=(2, 501))'))
+
+    def test_summarize_2d_dtype(self):
+        A = np.arange(1002, dtype='i2').reshape(2, 501)
+        strA = '[[   0    1    2 ...  498  499  500]\n' \
+               ' [ 501  502  503 ...  999 1000 1001]]'
+        assert_equal(str(A), strA)
+
+        reprA = ('array([[   0,    1,    2, ...,  498,  499,  500],\n'
+                 '       [ 501,  502,  503, ...,  999, 1000, 1001]],\n'
+                 '      shape=(2, 501), dtype=int16)')
         assert_equal(repr(A), reprA)
 
     def test_summarize_structure(self):
@@ -1040,7 +1063,7 @@ def test_edgeitems(self):
 
                    [[18, ..., 20],
                     ...,
-                    [24, ..., 26]]])""")
+                    [24, ..., 26]]], shape=(3, 3, 3))""")
         )
 
         b = np.zeros((3, 3, 1, 1))
@@ -1061,40 +1084,37 @@ def test_edgeitems(self):
 
                     ...,
 
-                    [[0.]]]])""")
+                    [[0.]]]], shape=(3, 3, 1, 1))""")
         )
 
         # 1.13 had extra trailing spaces, and was missing newlines
-        np.set_printoptions(legacy='1.13')
-
-        assert_equal(
-            repr(a),
-            textwrap.dedent("""\
-            array([[[ 0, ...,  2],
-                    ..., 
-                    [ 6, ...,  8]],
-
-                   ..., 
-                   [[18, ..., 20],
-                    ..., 
-                    [24, ..., 26]]])""")
-        )
-
-        assert_equal(
-            repr(b),
-            textwrap.dedent("""\
-            array([[[[ 0.]],
-
-                    ..., 
-                    [[ 0.]]],
-
-
-                   ..., 
-                   [[[ 0.]],
-
-                    ..., 
-                    [[ 0.]]]])""")
-        )
+        try:
+            np.set_printoptions(legacy='1.13')
+            assert_equal(repr(a), (
+                "array([[[ 0, ...,  2],\n"
+                "        ..., \n"
+                "        [ 6, ...,  8]],\n"
+                "\n"
+                "       ..., \n"
+                "       [[18, ..., 20],\n"
+                "        ..., \n"
+                "        [24, ..., 26]]])")
+            )
+            assert_equal(repr(b), (
+                "array([[[[ 0.]],\n"
+                "\n"
+                "        ..., \n"
+                "        [[ 0.]]],\n"
+                "\n"
+                "\n"
+                "       ..., \n"
+                "       [[[ 0.]],\n"
+                "\n"
+                "        ..., \n"
+                "        [[ 0.]]]])")
+            )
+        finally:
+            np.set_printoptions(legacy=False)
 
     def test_edgeitems_structured(self):
         np.set_printoptions(edgeitems=1, threshold=1)
