Depth-limited printing of types in stacktraces

base/show.jl

@@ -2462,6 +2462,290 @@ function show_signature_function(io::IO, @nospecialize(ft), demangle=false, farg
     nothing
 end
 
+"""
+    TypeTreesIO
+
+A module defining an IO subtype that constructs trees that mirror printing of a parametric types.
+"""
+module TypeTreesIO
+
+    export TypeTreeIO, TypeTreeNode
+
+    const delims = (('(', ')'), ('{', '}'))
+
+    mutable struct TypeTreeNode
+        name::String
+        parent::Union{Nothing,TypeTreeNode}
+        delimidx::Int8          # 1 or 2 for delims[delimidx], 0 for not assigned
+        children::Union{Nothing,Vector{TypeTreeNode}}
+    end
+    TypeTreeNode(name::AbstractString="", parent=nothing) = TypeTreeNode(name, parent, 0, nothing)
+
+    mutable struct TypeTreeBundle
+        body::TypeTreeNode                  # DataType
+        vars::Union{Nothing,TypeTreeNode}   # TypeVars
+    end
+    TypeTreeBundle(node::TypeTreeNode) = TypeTreeBundle(node, nothing)
+
+    """
+        TypeTreeIO() → io
+
+    Create an IO object to which you can print type objects or natural signatures.
+    Afterwards, `io.tree` will contain a tree representation of the printed type.
+
+    # Examples
+
+    ```jldoctest
+    julia> io = TypeTreeIO();
+
+    julia> print(io, Tuple{Int,Float64});
+
+    julia> io.tree.body.name
+    "Tuple"
+
+    julia> io.tree.body.children[1].name
+    "$Int"
+
+    julia> String(take!(io))
+    "Tuple{$Int, Float64}"
+    ```
+
+    # Extended help
+
+    In addition to printing a type directly to an `io::TypeTreeIO`, you can also
+    assemble it manually if you follow a few precautions:
+
+        - any `where` statement must be printed as `print(io, " where ")` or
+          `print(io, " where {")`. The `where` string argument may not have any
+          additional characters. Note the bracketing spaces.
+
+    """
+    mutable struct TypeTreeIO <: IO    # TODO?: abstract type TextIO <: IO end for text-only printing
+        io::Union{IOBuffer,IOContext{IOBuffer}}
+        tree::TypeTreeBundle
+        cursor::TypeTreeNode   # current position in the tree
+    end
+    function TypeTreeIO(io=IOBuffer())
+        root = TypeTreeNode()
+        return TypeTreeIO(io, TypeTreeBundle(root), root)
+    end
+
+    getttio(io::TypeTreeIO) = io
+    getttio(ioctx::IOContext{TypeTreeIO}) = ioctx.io
+    getio(io::TypeTreeIO) = io.io
+    getio(ioctx::IOContext{TypeTreeIO}) = getio(getttio(ioctx))
+
+    ## IO interface
+
+    function Base.flush(io::TypeTreeIO)
+        str = String(take!(io.io))
+        if !isempty(str)
+            curs = io.cursor
+            if curs.children === nothing
+                curs.children = TypeTreeNode[]
+            end
+            push!(curs.children, TypeTreeNode(str, curs))
+        end
+        return
+    end
+    Base.iswritable(::TypeTreeIO) = true
+
+    function Base.unsafe_write(io::TypeTreeIO, p::Ptr{UInt8}, nb::UInt)
+        str = String(unsafe_wrap(Array, p, (Int(nb),)))
+        if startswith(str, " where ")
+            @assert io.tree.vars === nothing
+            io.cursor = io.tree.vars = TypeTreeNode(" where ")
+            endswith(str, '{') && (io.cursor.delimidx = 2)
+            io.cursor.children = TypeTreeNode[]
+            return nb
+        end
+        for c in str
+            write(io, c)
+        end
+        return nb
+    end
+    Base.unsafe_write(ioctx::IOContext{TypeTreeIO}, p::Ptr{UInt8}, nb::UInt) = unsafe_write(getttio(ioctx), p, nb)
+
+    Base.get(treeio::TypeTreeIO, key, default) = get(treeio.io, key, default)
+
+    function Base.write(treeio::TypeTreeIO, c::Char)
+        curs = treeio.cursor
+        if c == '{' || (c == '(' && isempty(curs.name) && curs.parent === nothing)
+            str = String(take!(getio(treeio)))
+            if isempty(curs.name)
+                if curs.children === nothing
+                    curs.children = TypeTreeNode[]
+                end
+                curs.name = str
+                curs.delimidx = c == '(' ? 1 : 2
+            else
+                # We're dropping in depth
+                newcurs = TypeTreeNode(str, curs)
+                newcurs.delimidx = c == '(' ? 1 : 2
+                if curs.children === nothing
+                    curs.children = TypeTreeNode[]
+                end
+                push!(curs.children, newcurs)
+                treeio.cursor = newcurs
+            end
+        elseif c ∈ (',', '}') || (c == ')' && curs.delimidx == 1)
+            str = String(take!(getio(treeio)))
+            if !isempty(str)
+                if curs.children === nothing
+                    curs.children = TypeTreeNode[]
+                end
+                push!(curs.children, TypeTreeNode(str, curs))
+            else
+                if curs.children === nothing
+                    curs.children = TypeTreeNode[]
+                else
+                    p = curs.parent
+                    if p !== nothing
+                        treeio.cursor = p
+                    end
+                end
+            end
+        elseif c != ' '
+            print(treeio.io, c)
+        end
+        return textwidth(c)
+    end
+    Base.write(treeioctx::IOContext{TypeTreeIO}, c::Char) = write(getttio(treeioctx), c)
+
+
+    ## Printing the tree with constraints on width and/or depth
+
+    const truncchar = "…"
+    const per_param = ", "
+
+    function Base.take!(io::TypeTreeIO)
+        flush(io)
+        str = sprint(show, io.tree)
+        io.tree = TypeTreeBundle(TypeTreeNode())
+        io.cursor = io.tree.body
+        return codeunits(str)
+    end
+
+    function Base.show(io::IO, bundle::TypeTreeBundle)
+        depth = get(io, :type_depth, nothing)::Union{Int,Nothing}
+        if depth === nothing
+            maxdepth = get(io, :type_maxdepth, typemax(Int))::Int
+            maxwidth = get(io, :type_maxwidth, typemax(Int))::Int
+            depth = choose_depth(bundle, maxdepth, maxwidth)
+        end
+        _print(io, bundle.body, 1, depth, true)
+        vars = bundle.vars
+        if vars !== nothing
+            children = vars.children
+            if children !== nothing && length(children) == 1
+                vars = children[1]
+                print(io, " where ")
+            end
+            _print(io, vars, 1, depth, false)
+        end
+    end
+
+    function _print(io::IO, node::TypeTreeNode, thisdepth, maxdepth, isbody::Bool)
+        iscall = thisdepth==1 && isbody
+        isargs = thisdepth==2 && isbody
+        if isargs
+            sname = split(node.name, "::")
+            if length(sname) == 2
+                Base.print_within_stacktrace(io, sname[1]; color=:light_black)
+                print(io, "::")
+                print(io, sname[2])   # the top-level type is printed in :normal
+            else
+                print(io, node.name)
+            end
+        else
+            addparens = iscall && occursin("::", node.name)
+            addparens && print(io, '(')
+            Base.print_within_stacktrace(io, node.name; bold=iscall)
+            addparens && print(io, ')')
+        end
+        children = node.children
+        if children !== nothing
+            if isargs && get(io, :backtrace, false)::Bool
+                Base.with_output_color(:light_black, io) do io
+                    _print_children(io, children, node.delimidx, thisdepth, maxdepth, isbody)
+                end
+            else
+                _print_children(io, children, node.delimidx, thisdepth, maxdepth, isbody)
+            end
+        end
+    end
+
+    function _print_children(io::IO, children::Vector{TypeTreeNode}, delimidx, thisdepth, maxdepth, isbody)
+        bold = thisdepth==1 && isbody
+        iszero(delimidx) || Base.print_within_stacktrace(io, delims[delimidx][1]; bold)
+        if thisdepth >= maxdepth && !isempty(children)
+            print(io, truncchar)
+        else
+            n = lastindex(children)
+            for (i, child) in pairs(children)
+                _print(io, child, thisdepth+1, maxdepth, isbody)
+                i < n && print(io, per_param)
+            end
+        end
+        iszero(delimidx) || Base.print_within_stacktrace(io, delims[delimidx][end]; bold)
+        return
+    end
+
+    function choose_depth(wd::Vector{Int}, wtrunc::Vector{Int}, maxdepth::Int, maxwidth::Int)
+        wsum, depth = 0, 1
+        while depth <= maxdepth && depth <= length(wd)
+            wsum += wd[depth]
+            if wsum + wtrunc[depth] > maxwidth
+                return depth - 1
+            end
+            depth += 1
+        end
+        return depth - 1
+    end
+    choose_depth(bundle::TypeTreeBundle, maxdepth::Int, maxwidth::Int) =
+        choose_depth(width_by_depth(bundle)..., maxdepth, maxwidth)
+
+    """
+        width_by_depth(node) → wd, wtrunc
+
+    Compute the number of characters `wd[depth]` needed to print at each `depth` within the tree.
+    Also compute the number of additional characters `wtrunc[depth]` needed if one truncates the tree at `depth`.
+    """
+    function width_by_depth(bundle::TypeTreeBundle)
+        wd, wtrunc = Int[], Int[]
+        width_by_depth!(wd, wtrunc, bundle.body, 1)
+        if bundle.vars !== nothing
+            width_by_depth!(wd, wtrunc, bundle.vars, 1)
+        end
+        return wd, wtrunc
+    end
+
+    function width_by_depth!(wd, wtrunc, node::TypeTreeNode, depth)
+        if depth > length(wd)
+            push!(wd, 0)
+            push!(wtrunc, 0)
+        end
+        wd[depth] += length(node.name)
+        childs = node.children
+        if childs !== nothing
+            delimidx = node.delimidx
+            if !iszero(delimidx)
+                wd[depth] += length(delims[node.delimidx])
+            end
+            wtrunc[depth] += length(truncchar)
+            for child in childs
+                width_by_depth!(wd, wtrunc, child, depth+1)
+            end
+            if lastindex(wd) > depth
+                wd[depth+1] += length(per_param) * (length(childs) - 1)
+            end
+        end
+        return wd, wtrunc
+    end
+
+end
+using .TypeTreesIO
+
 function print_within_stacktrace(io, s...; color=:normal, bold=false)
     if get(io, :backtrace, false)::Bool
         printstyled(io, s...; color, bold)
@@ -2470,15 +2754,24 @@ function print_within_stacktrace(io, s...; color=:normal, bold=false)
     end
 end
 
-function show_tuple_as_call(io::IO, name::Symbol, sig::Type;
+function show_tuple_as_call(out::IO, name::Symbol, @nospecialize(sig::Type);
                             demangle=false, kwargs=nothing, argnames=nothing,
                             qualified=false, hasfirst=true)
     # print a method signature tuple for a lambda definition
     if sig === Tuple
-        print(io, demangle ? demangle_function_name(name) : name, "(...)")
+        print(out, demangle ? demangle_function_name(name) : name, "(...)")
         return
     end
     tv = Any[]
+    if get(out, :limit, false)::Bool
+        io = TypeTreeIO()
+        if isa(out, IOContext)
+            io = IOContext(io, out)
+            io = IOContext(io, :color => false)   # avoid ANSI escapes during tree construction (fix in `show`)
+        end
+    else
+        io = out
+    end
     env_io = io
     while isa(sig, UnionAll)
         push!(tv, sig.var)
@@ -2516,6 +2809,10 @@ function show_tuple_as_call(io::IO, name::Symbol, sig::Type;
     end
     print_within_stacktrace(io, ")", bold=true)
     show_method_params(io, tv)
+    if get(out, :limit, false)::Bool
+        sz = get(out, :displaysize, (typemax(Int), typemax(Int)))::Tuple{Int, Int}
+        show(IOContext(out, :type_maxwidth => sz[2]), TypeTreesIO.getttio(io).tree)
+    end
     nothing
 end
 

base/stacktraces.jl

@@ -222,6 +222,11 @@ function show_spec_linfo(io::IO, frame::StackFrame)
     elseif linfo isa MethodInstance
         def = linfo.def
         if isa(def, Method)
+            if get(io, :limit, false)::Bool
+                if !haskey(io, :displaysize)
+                    io = IOContext(io, :displaysize => displaysize(io))
+                end
+            end
             sig = linfo.specTypes
             argnames = Base.method_argnames(def)
             argnames = replace(argnames, :var"#unused#" => :var"")