Customize show for EinExpr, SizedEinExpr

bsc-quantic · Jul 25, 2024 · fbf8847 · fbf8847
1 parent b2b60ee
commit fbf8847
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/src/EinExpr.jl b/src/EinExpr.jl
@@ -17,6 +17,8 @@ EinExpr(head::NTuple, args::NTuple) = EinExpr(collect(head), collect(args))
 
 Base.copy(expr::EinExpr) = EinExpr(copy(expr.head), copy(expr.args))
 
+Base.show(io::IO, path::EinExpr) = print_tree((io, node) -> print(io, head(node)), io, path)
+
 """
     head(path::EinExpr)
 

diff --git a/src/SizedEinExpr.jl b/src/SizedEinExpr.jl
@@ -13,6 +13,16 @@ end
 
 EinExpr(path::Vector{L}, size::Dict{L}) where {L} = SizedEinExpr(EinExpr(path), size)
 
+function Base.show(io::IO, path::SizedEinExpr)
+    print_tree(io, path.path) do io, node
+        print(io, head(node))
+        print(
+            io,
+            " " * join(["$(flops(node, path.size)) flops", "$(length(SizedEinExpr(node, path.size))) elems"], ", "),
+        )
+    end
+end
+
 head(sexpr::SizedEinExpr) = head(sexpr.path)
 
 """
-Original file line number
+Diff line change
@@ Expand Up @@
     Base.copy(expr::EinExpr) = EinExpr(copy(expr.head), copy(expr.args))
+    Base.show(io::IO, path::EinExpr) = print_tree((io, node) -> print(io, head(node)), io, path)
     """
         head(path::EinExpr)
@@ Expand Down @@