further tweak tensor and block show (#322)

Author: Jutho

src/tensors/abstracttensor.jl

@@ -652,22 +652,22 @@ function Base.show(io::IO, mime::MIME"text/plain", t::AbstractTensorMap)
     # 1) show summary: typically d₁×d₂×… ← d₃×d₄×… $(typeof(t))
     summary(io, t)
 
-    # case without `\n`:
-    if get(io, :compact, true)
+    if get(io, :compact, false)
+        # case without `\n`:
         print(io, "(…, ")
         show(io, mime, space(t))
         print(io, ')')
-        return nothing
+    else
+        # case with `\n`
+        # 2) show spaces
+        println(io, ':')
+        println(io, " codomain: ", codomain(t))
+        println(io, " domain: ", domain(t))
+        # 3) show data
+        println(io, " blocks: ")
+        (numlines, numcols) = get(io, :displaysize, displaysize(io))
+        newio = IOContext(io, :displaysize => (numlines - 4, numcols))
+        show_blocks(newio, mime, blocks(t))
     end
-
-    # case with `\n`
-    # 2) show spaces
-    println(io, ':')
-    println(io, " codomain: ", codomain(t))
-    println(io, " domain: ", domain(t))
-
-    # 3) [optional]: show data
-    println(io, "\n\n blocks: ")
-    show_blocks(io, mime, blocks(t))
     return nothing
 end

src/tensors/blockiterator.jl

@@ -45,13 +45,40 @@ function foreachblock(f, t::AbstractTensorMap; scheduler = nothing)
     return nothing
 end
 
-function show_blocks(io, mime::MIME"text/plain", iter)
-    first = true
-    for (c, b) in iter
-        first || print(io, "\n\n")
-        print(io, " * ", c, " => ")
-        show(io, mime, b)
-        first = false
+function show_blocks(io, mime::MIME"text/plain", iter; maytruncate::Bool = true)
+    if maytruncate && get(io, :limit, false)
+        numlinesleft, numcols = get(io, :displaysize, displaysize(io))::Tuple{Int, Int}
+        numlinesleft -= 2 # lines of headers have already been subtracted, but not the 2 spare lines for old and new prompts
+        minlinesperblock = 7 # aim to have at least this many lines per printed block (= 5 lines for the actual matrix)
+        minnumberofblocks = min(3, length(iter)) # aim to show at least this many blocks
+        truncateblocks = sum(cb -> min(size(cb[2], 1) + 2, minlinesperblock), iter; init = 0) > numlinesleft
+        maxnumlinesperblock = max(div(numlinesleft - 2 * truncateblocks, minnumberofblocks), minlinesperblock)
+        # aim to show at least minnumberofblocks, but not if this means that there would be less than minlinesperblock
+        # deduct two lines for a truncation message (and newline) if needed
+        for (n, (c, b)) in enumerate(iter)
+            n == 1 || print(io, "\n\n")
+            numlinesneeded = min(size(b, 1) + 2, maxnumlinesperblock)
+            if numlinesleft >= numlinesneeded + 2 * truncateblocks
+                # we can still print at least this block, and have two lines for
+                # the truncation message (and its newline) if it is required
+                print(io, " * ", c, " => ")
+                newio = IOContext(io, :displaysize => (maxnumlinesperblock - 1 + 3, numcols))
+                # subtract 1 line for the newline, but add 3 because of how matrices are printed
+                show(newio, mime, b)
+                numlinesleft -= numlinesneeded
+            else
+                print(io, " * ", "  \u2026   [output of ", length(iter) - n + 1, " more block(s) truncated]")
+                break
+            end
+        end
+    else
+        first = true
+        for (c, b) in iter
+            first || print(io, "\n\n")
+            print(io, " * ", c, " => ")
+            show(io, mime, b)
+            first = false
+        end
     end
     return nothing
 end
@@ -73,7 +100,9 @@ end
 function Base.show(io::IO, mime::MIME"text/plain", b::BlockIterator)
     summary(io, b)
     println(io, ":")
-    show_blocks(io, mime, b)
+    (numlines, numcols) = get(io, :displaysize, displaysize(io))::Tuple{Int, Int}
+    newio = IOContext(io, :displaysize => (numlines - 1, numcols))
+    show_blocks(newio, mime, b; maytruncate = false)
     return nothing
 end
 

test/tensors/tensors.jl

@@ -624,3 +624,52 @@ end
         end
     end
 end
+
+@timedtestset "show tensors" begin
+    for V in (ℂ^2, Z2Space(0 => 2, 1 => 2), SU2Space(0 => 2, 1 => 2))
+        t1 = ones(Float32, V ⊗ V, V)
+        t2 = randn(ComplexF64, V ⊗ V ⊗ V)
+        # test unlimited output
+        for t in (t1, t2, t1', t2')
+            output = IOBuffer()
+            summary(output, t)
+            print(output, ":\n codomain: ")
+            show(output, MIME("text/plain"), codomain(t))
+            print(output, "\n domain: ")
+            show(output, MIME("text/plain"), domain(t))
+            print(output, "\n blocks: \n")
+            first = true
+            for (c, b) in blocks(t)
+                first || print(output, "\n\n")
+                print(output, " * ")
+                show(output, MIME("text/plain"), c)
+                print(output, " => ")
+                show(output, MIME("text/plain"), b)
+                first = false
+            end
+            outputstr = String(take!(output))
+            @test outputstr == sprint(show, MIME("text/plain"), t)
+        end
+
+        # test limited output with a single block
+        t = randn(Float64, V ⊗ V, V)' # we know there is a single space in the codomain, so that blocks have 2 rows
+        output = IOBuffer()
+        summary(output, t)
+        print(output, ":\n codomain: ")
+        show(output, MIME("text/plain"), codomain(t))
+        print(output, "\n domain: ")
+        show(output, MIME("text/plain"), domain(t))
+        print(output, "\n blocks: \n")
+        c = unit(sectortype(t))
+        b = block(t, c)
+        print(output, " * ")
+        show(output, MIME("text/plain"), c)
+        print(output, " => ")
+        show(output, MIME("text/plain"), b)
+        if length(blocks(t)) > 1
+            print(output, "\n\n *   …   [output of 1 more block(s) truncated]")
+        end
+        outputstr = String(take!(output))
+        @test outputstr == sprint(show, MIME("text/plain"), t; context = (:limit => true, :displaysize => (12, 100)))
+    end
+end