Fix maxdepth for print_tree and add tests (#42)

Author: jonalm

src/AbstractTrees.jl

@@ -98,10 +98,13 @@ struct TreeCharSet
     terminator
     skip
     dash
+    trunc
 end
 
 # Default charset
-TreeCharSet() = TreeCharSet('├','└','│','─')
+TreeCharSet() = TreeCharSet('├','└','│','─','⋮')
+TreeCharSet(mid, term, skip, dash) = TreeCharSet(mid, term, skip, dash, '⋮')
+
 
 function print_prefix(io, depth, charset, active_levels)
     for current_depth in 0:(depth-1)
@@ -113,8 +116,9 @@ function print_prefix(io, depth, charset, active_levels)
     end
 end
 
-function _print_tree(printnode::Function, io::IO, tree, maxdepth = 5; depth = 0, active_levels = Int[],
-                     charset = TreeCharSet(), withinds = false, inds = [], from = nothing, to = nothing, roottree = tree)
+function _print_tree(printnode::Function, io::IO, tree, maxdepth = 5; indicate_truncation = true,
+                     depth = 0, active_levels = Int[], charset = TreeCharSet(), withinds = false,
+                     inds = [], from = nothing, to = nothing, roottree = tree)
     nodebuf = IOBuffer()
     isa(io, IOContext) && (nodebuf = IOContext(nodebuf, io))
     if withinds
@@ -132,23 +136,31 @@ function _print_tree(printnode::Function, io::IO, tree, maxdepth = 5; depth = 0,
     c = isa(treekind(roottree), IndexedTree) ?
         childindices(roottree, tree) : children(roottree, tree)
     if c !== ()
-        s = Iterators.Stateful(from === nothing ? pairs(c) : Iterators.Rest(pairs(c), from))
-        while !isempty(s)
-            ind, child = popfirst!(s)
-            ind === to && break
-            active = false
-            child_active_levels = active_levels
-            print_prefix(io, depth, charset, active_levels)
-            if isempty(s)
-                print(io, charset.terminator)
-            else
-                print(io, charset.mid)
-                child_active_levels = push!(copy(active_levels), depth)
+        if depth < maxdepth
+            s = Iterators.Stateful(from === nothing ? pairs(c) : Iterators.Rest(pairs(c), from))
+            while !isempty(s)
+                ind, child = popfirst!(s)
+                ind === to && break
+                active = false
+                child_active_levels = active_levels
+                print_prefix(io, depth, charset, active_levels)
+                if isempty(s)
+                    print(io, charset.terminator)
+                else
+                    print(io, charset.mid)
+                    child_active_levels = push!(copy(active_levels), depth)
+                end
+                print(io, charset.dash, ' ')
+                print_tree(printnode, io, child, maxdepth;
+                indicate_truncation=indicate_truncation, depth = depth + 1,
+                active_levels = child_active_levels, charset = charset, withinds=withinds,
+                inds = withinds ? [inds; ind] : [], roottree = roottree)
             end
-            print(io, charset.dash, ' ')
-            print_tree(printnode, io, child; depth = depth + 1,
-              active_levels = child_active_levels, charset = charset, withinds=withinds,
-              inds = withinds ? [inds; ind] : [], roottree = roottree)
+        elseif indicate_truncation
+            print_prefix(io, depth, charset, active_levels)
+            println(io, charset.trunc)
+            print_prefix(io, depth, charset, active_levels)
+            println(io)
         end
     end
 end
@@ -164,23 +176,31 @@ print_tree(tree, args...; kwargs...) = print_tree(stdout::IO, tree, args...; kwa
 # Usage
 Prints an ASCII formatted representation of the `tree` to the given `io` object.
 By default all children will be printed up to a maximum level of 5, though this
-value can be overriden by the `maxdepth` parameter. The charset to use in
+value can be overriden by the `maxdepth` parameter. Nodes that are truncated are
+indicated by a vertical ellipsis below the truncated node, this indication can be
+turned off by providing `indicate_truncation=false` as a kwarg. The charset to use in
 printing can be customized using the `charset` keyword argument.
 You can control the printing of individual nodes by passing a function `f(io, node)`;
 the default is [`AbstractTrees.printnode`](@ref).
 
 # Examples
 ```julia
-julia> print_tree(STDOUT,Dict("a"=>"b","b"=>['c','d']))
+julia> print_tree(stdout, Dict("a"=>"b","b"=>['c','d']))
 Dict{String,Any}("b"=>['c','d'],"a"=>"b")
 ├─ b
 │  ├─ c
 │  └─ d
 └─ a
    └─ b
 
-julia> print_tree(STDOUT,Dict("a"=>"b","b"=>['c','d']);
-        charset = TreeCharSet('+','\\','|',"--"))
+julia> print_tree(stdout, '0'=>'1'=>'2'=>'3', 2)
+'0'
+└─ '1'
+    └─ '2'
+        ⋮
+
+julia> print_tree(stdout, Dict("a"=>"b","b"=>['c','d']);
+        charset = TreeCharSet('+','\\','|',"--","⋮"))
 Dict{String,Any}("b"=>['c','d'],"a"=>"b")
 +-- b
 |   +-- c

test/runtests.jl

@@ -138,3 +138,85 @@ mktemp() do filename, io
     end
 end
 end  # @testset "Examples"
+
+struct Num{I} end
+Num(I::Int) = Num{I}()
+Base.show(io::IO, ::Num{I}) where {I} = print(io, I)
+AbstractTrees.children(x::Num{I}) where {I} = (Num(I+1), Num(I+1))
+
+struct SingleChildInfiniteDepth end
+AbstractTrees.children(::SingleChildInfiniteDepth) = (SingleChildInfiniteDepth(),)
+
+@testset "Test print_tree truncation" begin
+
+# test that `print_tree(headnode, maxdepth)` truncates the output at right depth
+# julia > print_tree(Num(0), 3)
+# 0
+# ├─ 1
+# │  ├─ 2
+# │  │  ├─ 3
+# │  │  └─ 3
+# │  └─ 2
+# │     ├─ 3
+# │     └─ 3
+# └─ 1
+#    ├─ 2
+#    │  ├─ 3
+#    │  └─ 3
+#    └─ 2
+#       ├─ 3
+#       └─ 3
+#
+
+for maxdepth in [3,5,8]
+    buffer = IOBuffer()
+    print_tree(buffer, Num(0), maxdepth)
+    ptxt = String(take!(buffer))
+    n1  = sum([1 for c in ptxt if c=="$(maxdepth-1)"[1]])
+    n2  = sum([1 for c in ptxt if c=="$maxdepth"[1]])
+    n3  = sum([1 for c in ptxt if c=="$(maxdepth+1)"[1]])
+    @test n1==2^(maxdepth-1)
+    @test n2==2^maxdepth
+    @test n3==0
+end
+
+# test that `print_tree(headnode)` prints truncation characters under each
+# node at the default maxdepth level = 5
+truncation_char = AbstractTrees.TreeCharSet().trunc
+buffer = IOBuffer()
+print_tree(buffer, Num(0))
+ptxt = String(take!(buffer))
+n1  = sum([1 for c in ptxt if c=='5'])
+n2  = sum([1 for c in ptxt if c=='6'])
+n3  = sum([1 for c in ptxt if c==truncation_char])
+@test n1==2^5
+@test n2==0
+@test n3==2^5
+lines = split(ptxt, '\n')
+for i in 1:length(lines)
+    if ~isempty(lines[i]) && lines[i][end] == '5'
+        @test lines[i+1][end] == truncation_char
+    end
+end
+
+# test correct number of lines printed 1
+buffer = IOBuffer()
+print_tree(buffer, SingleChildInfiniteDepth())
+ptxt = String(take!(buffer))
+numlines = sum([1 for c in split(ptxt, '\n') if ~isempty(strip(c))])
+@test numlines == 7 # 1 (head node) + 5 (default depth) + 1 (truncation char)
+
+# test correct number of lines printed 2
+buffer = IOBuffer()
+print_tree(buffer, SingleChildInfiniteDepth(), 3)
+ptxt = String(take!(buffer))
+numlines = sum([1 for c in split(ptxt, '\n') if ~isempty(strip(c))])
+@test numlines == 5 # 1 (head node) + 3 (depth) + 1 (truncation char)
+
+# test correct number of lines printed 3
+buffer = IOBuffer()
+print_tree(buffer, SingleChildInfiniteDepth(), 3, indicate_truncation=false)
+ptxt = String(take!(buffer))
+numlines = sum([1 for c in split(ptxt, '\n') if ~isempty(strip(c))])
+@test numlines == 4 # 1 (head node) + 3 (depth)
+end