tuple -> tuple[type1, ... typen] in tree_summary type row

ASEM000 · Mar 29, 2024 · 07af95f · 07af95f
1 parent 810983b
commit 07af95f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/sepes/_src/tree_pprint.py b/sepes/_src/tree_pprint.py
@@ -570,6 +570,21 @@ def _(node: Any) -> str:
         return tree_repr(ShapeDTypePP(shape, dtype))
 
 
+@tree_summary.def_type(tuple)
+def _(node: tuple) -> str:
+    # - output tuple[types,...] instead of just tuple in the type row.
+    # - usually this encounterd if the tree_summary depth is not inf
+    #   so the tree leaves could contain non-atomic types.
+    treelib = sepes._src.backend.treelib
+
+    one_level_types = treelib.map(
+        tree_summary.type_dispatcher,
+        node,
+        is_leaf=lambda x: False if id(x) == id(node) else True,
+    )
+    return f"tuple[{','.join(one_level_types)}]"
+
+
 if is_package_avaiable("jax"):
     # jax pretty printing extra handlers
     import jax

diff --git a/tests/test_pprint.py b/tests/test_pprint.py
@@ -141,7 +141,7 @@ def test_tree_summary():
     assert (
         tree_summary(r1, depth=1)
         # trunk-ignore(flake8/E501)
-        == "┌────┬────────────┬─────┬───────┐\n│Name│Type        │Count│Size   │\n├────┼────────────┼─────┼───────┤\n│.a  │int         │1    │       │\n├────┼────────────┼─────┼───────┤\n│.b  │str         │1    │       │\n├────┼────────────┼─────┼───────┤\n│.c  │float       │1    │       │\n├────┼────────────┼─────┼───────┤\n│.d  │str         │1    │       │\n├────┼────────────┼─────┼───────┤\n│.e  │list        │5    │       │\n├────┼────────────┼─────┼───────┤\n│.f  │set         │1    │       │\n├────┼────────────┼─────┼───────┤\n│.g  │dict        │27   │100.00B│\n├────┼────────────┼─────┼───────┤\n│.h  │f32[5,1]    │5    │20.00B │\n├────┼────────────┼─────┼───────┤\n│.i  │f32[1,6]    │6    │24.00B │\n├────┼────────────┼─────┼───────┤\n│.j  │f32[1,1,4,5]│20   │80.00B │\n├────┼────────────┼─────┼───────┤\n│.k  │tuple       │3    │       │\n├────┼────────────┼─────┼───────┤\n│.l  │a           │2    │       │\n├────┼────────────┼─────┼───────┤\n│.m  │f32[5,5]    │25   │100.00B│\n├────┼────────────┼─────┼───────┤\n│.n  │bool[]      │1    │1.00B  │\n├────┼────────────┼─────┼───────┤\n│.o  │c64[2]      │2    │16.00B │\n├────┼────────────┼─────┼───────┤\n│Σ   │Repr1       │101  │341.00B│\n└────┴────────────┴─────┴───────┘"
+        == '┌────┬──────────────────┬─────┬───────┐\n│Name│Type              │Count│Size   │\n├────┼──────────────────┼─────┼───────┤\n│.a  │int               │1    │       │\n├────┼──────────────────┼─────┼───────┤\n│.b  │str               │1    │       │\n├────┼──────────────────┼─────┼───────┤\n│.c  │float             │1    │       │\n├────┼──────────────────┼─────┼───────┤\n│.d  │str               │1    │       │\n├────┼──────────────────┼─────┼───────┤\n│.e  │list              │5    │       │\n├────┼──────────────────┼─────┼───────┤\n│.f  │set               │1    │       │\n├────┼──────────────────┼─────┼───────┤\n│.g  │dict              │27   │100.00B│\n├────┼──────────────────┼─────┼───────┤\n│.h  │f32[5,1]          │5    │20.00B │\n├────┼──────────────────┼─────┼───────┤\n│.i  │f32[1,6]          │6    │24.00B │\n├────┼──────────────────┼─────┼───────┤\n│.j  │f32[1,1,4,5]      │20   │80.00B │\n├────┼──────────────────┼─────┼───────┤\n│.k  │tuple[int,int,int]│3    │       │\n├────┼──────────────────┼─────┼───────┤\n│.l  │tuple[int,int]    │2    │       │\n├────┼──────────────────┼─────┼───────┤\n│.m  │f32[5,5]          │25   │100.00B│\n├────┼──────────────────┼─────┼───────┤\n│.n  │bool[]            │1    │1.00B  │\n├────┼──────────────────┼─────┼───────┤\n│.o  │c64[2]            │2    │16.00B │\n├────┼──────────────────┼─────┼───────┤\n│Σ   │Repr1             │101  │341.00B│\n└────┴──────────────────┴─────┴───────┘'
     )
 
     assert (