TST: Add test for the nditer debug_print

1 files changed, 85 insertions, 0 deletions

diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
index 94f61baca..536f812ee 100644
--- a/numpy/core/tests/test_nditer.py
+++ b/numpy/core/tests/test_nditer.py
@@ -1,6 +1,8 @@
 import sys
 import pytest
 
+import subprocess, textwrap, re
+
 import numpy as np
 import numpy.core._multiarray_tests as _multiarray_tests
 from numpy import array, arange, nditer, all
@@ -3014,3 +3016,86 @@ def test_partial_iteration_error(in_dtype, buf_dtype):
         it.iternext()
 
     assert count == sys.getrefcount(value)
+
+
+def test_debug_print():
+    """
+    Matches the expected output of a debug print with the actual output.
+    Note that the iterator dump should not be considered stable API,
+    this test is mainly to ensure the print does not crash.
+
+    Currently uses a subprocess to avoid dealing with the C level `printf`s.
+    """
+    # the expected output with all addresses and sizes stripped (they vary
+    # and/or are platform dependend).
+    expected = textwrap.dedent("""
+    ------ BEGIN ITERATOR DUMP ------
+    | Iterator Address:
+    | ItFlags: BUFFER REDUCE REUSE_REDUCE_LOOPS
+    | NDim: 2
+    | NOp: 2
+    | IterSize: 50
+    | IterStart: 0
+    | IterEnd: 50
+    | IterIndex: 0
+    | Iterator SizeOf:
+    | BufferData SizeOf:
+    | AxisData SizeOf:
+    |
+    | Perm: 0 1
+    | DTypes:
+    | DTypes: dtype('float64') dtype('int32')
+    | InitDataPtrs:
+    | BaseOffsets: 0 0
+    | Operands:
+    | Operand DTypes: dtype('int64') dtype('float64')
+    | OpItFlags:
+    |   Flags[0]: READ CAST ALIGNED
+    |   Flags[1]: READ WRITE CAST ALIGNED REDUCE
+    |
+    | BufferData:
+    |   BufferSize: 50
+    |   Size: 5
+    |   BufIterEnd: 5
+    |   REDUCE Pos: 0
+    |   REDUCE OuterSize: 10
+    |   REDUCE OuterDim: 1
+    |   Strides: 8 4
+    |   Ptrs:
+    |   REDUCE Outer Strides: 40 0
+    |   REDUCE Outer Ptrs:
+    |   ReadTransferFn:
+    |   ReadTransferData: (nil) (nil)
+    |   WriteTransferFn: (nil)
+    |   WriteTransferData: (nil) (nil)
+    |   Buffers:
+    |
+    | AxisData[0]:
+    |   Shape: 5
+    |   Index: 0
+    |   Strides: 16 8
+    |   Ptrs:
+    | AxisData[1]:
+    |   Shape: 10
+    |   Index: 0
+    |   Strides: 80 0
+    |   Ptrs:
+    ------- END ITERATOR DUMP -------
+    """).strip()
+
+    code = textwrap.dedent("""
+        import numpy as np
+        arr1 = np.arange(100, dtype=np.int64).reshape(10, 10)[:, ::2]
+        arr2 = np.arange(5.)
+        it = np.nditer((arr1, arr2), op_dtypes=["d", "i4"], casting="unsafe",
+                       flags=["reduce_ok", "buffered"],
+                       op_flags=[["readonly"], ["readwrite"]])
+        it.debug_print()
+        """)
+    res = subprocess.check_output([sys.executable, "-c", code], text=True)
+    res = res.strip()
+
+    for res_line, expected_line in zip(res.splitlines(), expected.splitlines()):
+        # The actual output may have additional pointers listed that are
+        # stripped from the example output:
+        assert res_line.startswith(expected_line)