Merge pull request JuliaCollections#124 from FedeClaudi/master

Increase customization in `print_tree`

Author: oscardssmith

src/printing.jl

@@ -1,14 +1,16 @@
 """
-    print_tree(tree; kwargs...)
-    print_tree(io::IO, tree; kwargs...)
-    print_tree(f::Function, io::IO, tree; kwargs...)
+    print_tree(tree; kw...)
+    print_tree(io::IO, tree; kw...)
+    print_tree(f::Function, g::Function, io::IO, tree; kw...)
 
 Print a text representation of `tree` to the given `io` object.
 
 # Arguments
 
 * `f::Function` - custom implementation of [`printnode`](@ref) to use. Should have the
-  signature `f(io::IO, node)`.
+  signature `f(io::IO, node; kw...)`.
+* `g::Function` -  custom implementation of [`print_child_key`](@ref) to use. Should have the
+signature `g(io::IO, key;)`.
 * `io::IO` - IO stream to write to.
 * `tree` - tree to print.
 * `maxdepth::Integer = 5` - truncate printing of subtrees at this depth.
@@ -82,7 +84,7 @@ Print a compact representation of a single node.  By default, this prints `nodev
 **OPTIONAL**: This can be extended for custom types and controls how nodes are shown
 in [`print_tree`](@ref).
 """
-printnode(io::IO, node) = show(IOContext(io, :compact => true, :limit => true), nodevalue(node))
+printnode(io::IO, node; kw...) = show(IOContext(io, :compact => true, :limit => true), nodevalue(node))
 
 
 """
@@ -185,17 +187,18 @@ print_child_key(io::IO, key::CartesianIndex) = show(io, Tuple(key))
 
 branchwidth(cs::TreeCharSet) = sum(textwidth.((cs.mid, cs.dash)))
 
-function print_tree(printnode::Function, io::IO, node;
+function print_tree(printnode::Function, print_child_key::Function, io::IO, node;
                     maxdepth::Integer=5,
                     indicate_truncation::Bool=true,
                     charset::TreeCharSet=TreeCharSet(),
                     printkeys::Union{Bool,Nothing}=nothing,
                     depth::Integer=0,
                     prefix::AbstractString="",
+                    kw...
                    )
     # Get node representation as string
     buf = IOBuffer()
-    printnode(IOContext(buf, io), node)
+    printnode(IOContext(buf, io), node; kw...)
     str = String(take!(buf))
 
     # Copy buffer to output, prepending prefix to each line
@@ -259,14 +262,14 @@ function print_tree(printnode::Function, io::IO, node;
             child_prefix *= " " ^ (textwidth(key_str) + textwidth(charset.pair))
         end
 
-        print_tree(printnode, io, child;
+        print_tree(printnode, print_child_key, io, child;
                    maxdepth=maxdepth, indicate_truncation=indicate_truncation, charset=charset,
                    printkeys=printkeys, depth=depth+1, prefix=child_prefix
                   )
     end
 end
 
-print_tree(io::IO, node; kw...) = print_tree(printnode, io, node; kw...)
+print_tree(io::IO, node; kw...) = print_tree(printnode, print_child_key, io, node; kw...)
 print_tree(node; kw...) = print_tree(stdout, node; kw...)
 
 
@@ -278,10 +281,10 @@ Get the string result of calling [`print_tree`](@ref) with the supplied argument
 
 The `context` argument works as it does in `Base.repr`.
 """
-repr_tree(tree; context=nothing, kw...) = repr_tree(printnode, tree; context=nothing, kw...)
-function repr_tree(f, tree; context=nothing, kw...)
+repr_tree(tree; context=nothing, kw...) = repr_tree(printnode, print_child_key, tree; context=nothing, kw...)
+function repr_tree(f::Function, g::Function, tree; context=nothing, kw...)
     buf = IOBuffer()
     io = context === nothing ? buf : IOContext(buf, context)
-    print_tree(f, io, tree; kw...)
+    print_tree(f, g, io, tree; kw...)
     return String(take!(buf))
 end

test/printing.jl

@@ -190,14 +190,10 @@ end
                    """)
 
     # Test printnode override
-    str = repr_tree(tree) do io, s
-        if s isa BoxNode
-            print(io, s.s)
-        else
-            AbstractTrees.printnode(io, s)
-        end
-    end
+    _f(io, s) = s isa BoxNode ? print(io, s.s) : AbstractTrees.printnode(io, s)
+    _g(io, k) = AbstractTrees.print_child_key(io, k)
 
+    str = repr_tree(_f, _g, tree) 
     @test endswith(str, """
                    ├─ "foo"
                    ├─ bar