Shorten printing of Tangent and thunks (#564)

* don't print internal details of thunks

* don't print super-long primal types for Tangent

* add tests

* version

* fix Int32, comments

* fix tests on 1.5

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.15.1"
+version = "1.15.2"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/tangent_types/tangent.jl

@@ -75,7 +75,18 @@ Base.hash(a::Tangent, h::UInt) = Base.hash(backing(canonicalize(a)), h)
 
 function Base.show(io::IO, tangent::Tangent{P}) where {P}
     print(io, "Tangent{")
-    show(io, P)
+    str = sprint(show, P, context = io)
+    i = findfirst('{', str)
+    if isnothing(i)
+        print(io, str)
+    else  # for Tangent{T{A,B,C}}(stuff), print {A,B,C} in grey, and trim this part if longer than a line:
+        print(io, str[1:prevind(str, i)])
+        if length(str) < 80
+            printstyled(io, str[i:end], color=:light_black)
+        else
+           printstyled(io, str[i:prevind(str, 80)], "...", color=:light_black) 
+        end
+    end
     print(io, "}")
     if isempty(backing(tangent))
         print(io, "()")  # so it doesn't show `NamedTuple()`

src/tangent_types/thunks.jl

@@ -196,7 +196,13 @@ end
 
 function Base.show(io::IO, x::Thunk)
     print(io, "Thunk(")
-    show(io, x.f)
+    str = sprint(show, x.f, context = io)  # often this name is like "ChainRules.var"#1398#1403"{Matrix{Float64}, Matrix{Float64}}"
+    ind = findfirst("var\"#", str)
+    if isnothing(ind) || length(str) < 80
+        printstyled(io, str, color=:light_black)
+    else
+        printstyled(io, str[1:ind[5]], "...", color=:light_black)
+    end
     print(io, ")")
 end
 
@@ -223,7 +229,13 @@ unthunk(x::InplaceableThunk) = unthunk(x.val)
 
 function Base.show(io::IO, x::InplaceableThunk)
     print(io, "InplaceableThunk(")
-    show(io, x.add!)
+    str = sprint(show, x.add!, context = io)
+    ind = findfirst("var\"#", str)  # look for auto-generated function names, often with huge types
+    if isnothing(ind)
+        printstyled(io, str, color=:light_black)
+    else
+        printstyled(io, str[1:ind[5]], "...", color=:light_black)
+    end
     print(io, ", ")
     show(io, x.val)
     print(io, ")")

test/tangent_types/tangent.jl

@@ -380,4 +380,12 @@ end
         c = Tangent{typeof(nt)}(; a=NoTangent(), b=0.1)
         @test nt + c == (; a=1, b=2.1)
     end
+    
+    @testset "printing" begin
+        t5 = Tuple(rand(3))
+        nt3 = (x=t5, y=t5, z=nothing)
+        tang = ProjectTo(nt3)(nt3)  # moderately complicated Tangent
+        @test contains(sprint(show, tang), "...}(x = Tangent")  # gets shortened
+        @test contains(sprint(show, tang), sprint(show, tang.x))  # inner piece appears whole
+    end
 end

test/tangent_types/thunks.jl

@@ -190,4 +190,16 @@
         @test scal!(2, 2.0, v, 1) == scal!(2, @thunk(2.0), v, 1)
         @test_throws MutateThunkException LAPACK.trsyl!('C', 'C', m, m, @thunk(m))
     end
+    
+    @testset "printing" begin
+        @test !contains(sprint(show, @thunk 1+1), "...")  # short thunks not abbreviated
+        th = let x = rand(100)
+            @thunk x .+ x'
+        end
+        @test contains(sprint(show, th), "...")  # but long ones are
+        
+        @test contains(sprint(show, InplaceableThunk(mul!, th)), "mul!")  # named functions left in InplaceableThunk
+        str = sprint(show, InplaceableThunk(z -> z .+ ones(100), th))
+        @test length(findall("...", str)) == 2  # now both halves shortened
+    end
 end