fix display

sdpython · sdpython · Feb 21, 2024 · Feb 2, 2024 · Feb 2, 2024 · Feb 2, 2024
commit 463bc86b1f10c1c5e2cd1b5ee8ad3b9701f96f7a
diff --git a/_unittests/ut_reference/test_evaluator_yield.py b/_unittests/ut_reference/test_evaluator_yield.py
@@ -424,13 +424,13 @@ def test_distance_sequence_str(self):
         text = dc.to_str(s1, s2, align)
         self.assertIn("OUTPUT", text)
         expected = """
-            001=|INPUTfloat322x2ABCDA|INPUTfloat322x2ABCDA
-            002=|INPUTfloat322x2ABCDB|INPUTfloat322x2ABCDB
-            003~|INPUTfloat322x3ABCDX|INPUTfloat322x2ABCDX
-            004-|RESULTfloat322x2CEIOExpH|
-            005=|RESULTfloat322x2CEIOLinearRegressioY1|RESULTfloat322x2CEIOLinearRegressioY1
-            006~|RESULTfloat322x2CEIOAbsY|RESULTfloat322x3CEIPAbsZ
-            007~|OUTPUTfloat322x2CEIOY|OUTPUTfloat322x2CEIPY
+            001=|INPUTfloat322:2x2ABCDA|INPUTfloat322:2x2ABCDA
+            002=|INPUTfloat322:2x2ABCDB|INPUTfloat322:2x2ABCDB
+            003~|INPUTfloat322:2x3ABCDX|INPUTfloat322:2x2ABCDX
+            004-|RESULTfloat322:2x2CEIOExpH|
+            005=|RESULTfloat322:2x2CEIOLinearRegressioY1|RESULTfloat322:2x2CEIOLinearRegressioY1
+            006~|RESULTfloat322:2x2CEIOAbsY|RESULTfloat322:2x3CEIPAbsZ
+            007~|OUTPUTfloat322:2x2CEIOY|OUTPUTfloat322:2x2CEIPY
         """.replace(
             "            ", ""
         ).strip(

diff --git a/onnx_array_api/reference/evaluator_yield.py b/onnx_array_api/reference/evaluator_yield.py
@@ -29,6 +29,22 @@ def __repr__(self):
         return f"{self.__class__.__name__}.{self._name_}"
 
 
+def _dimension_to_str(d):
+    if isinstance(d, int):
+        return str(d)
+    try:
+        int(d)
+    except ValueError:
+        return d
+    return f"{d!r}"
+
+
+def _rank_to_str(shape):
+    if shape:
+        return f"{len(shape)}:"
+    return "  "
+
+
 @dataclass
 class ResultExecution:
     """
@@ -65,7 +81,13 @@ def __str__(self):
         els = [
             _align(self.kind._name_, 6),
             _align(str(dtype).replace("dtype(", "").replace(")", ""), 8),
-            _align("x".join("" if self.shape is None else map(str, self.shape)), 15),
+            _rank_to_str(self.shape)
+            + _align(
+                "x".join(
+                    "" if self.shape is None else map(_dimension_to_str, self.shape)
+                ),
+                18,
+            ),
             self.summary,
             _align(self.op_type or "", 15),
             self.name or "",