Handle nested progress bars (#13)

Author: tkf

Project.toml

@@ -4,10 +4,15 @@ authors = ["Chris Foster <chris42f@gmail.com>"]
 version = "0.1.0"
 
 [deps]
+LeftChildRightSiblingTrees = "1d6d02ad-be62-4b6b-8a6d-2f90e265016e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
+UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
+LeftChildRightSiblingTrees = "0.1.1"
+ProgressLogging = "0.1.1"
 julia = "1"
 
 [extras]

src/ProgressBar.jl

@@ -0,0 +1,49 @@
+Base.@kwdef mutable struct ProgressBar
+    fraction::Union{Float64,Nothing}
+    name::String
+    level::Int = 0
+    barglyphs::BarGlyphs = BarGlyphs()
+    tlast::Float64 = time()
+    tfirst::Float64 = time()
+    id::UUID
+    parentid::UUID
+end
+
+set_fraction!(bar::ProgressBar, ::Nothing) = bar
+function set_fraction!(bar::ProgressBar, fraction::Real)
+    bar.tlast = time()
+    bar.fraction = fraction
+    return bar
+end
+
+# This is how `ProgressMeter.printprogress` decides "ETA" vs "Time":
+ensure_done!(bar::ProgressBar) = bar.fraction = 1
+
+function eta_seconds(bar)
+    total = (bar.tlast - bar.tfirst) / something(bar.fraction, NaN)
+    return total - (time() - bar.tfirst)
+end
+
+function printprogress(io::IO, bar::ProgressBar)
+    if bar.name == ""
+        desc = "Progress: "
+    else
+        desc = bar.name
+        if !endswith(desc, " ")
+            desc *= " "
+        end
+    end
+
+    pad = "  "^bar.level
+    print(io, pad)
+
+    lines, columns = displaysize(io)
+    ProgressMeter.printprogress(
+        IOContext(io, :displaysize => (lines, max(1, columns - length(pad)))),
+        bar.barglyphs,
+        bar.tfirst,
+        desc,
+        something(bar.fraction, NaN),
+        eta_seconds(bar),
+    )
+end

src/ProgressMeter/ProgressMeter.jl

@@ -21,43 +21,43 @@ BarGlyphs() = BarGlyphs(
     '|',
 )
 
-struct ProgressBar
-    barglyphs::BarGlyphs
-    tfirst::Float64
-end
-
-ProgressBar(barglyphs = BarGlyphs()) = ProgressBar(barglyphs, time())
-
 """
-    printprogress(io::IO, p::ProgressBar, desc, progress::Real)
+    printprogress(io::IO, barglyphs::BarGlyphs, tfirst::Float64, desc, progress::Real)
 
-Print progress bar to `io` with setting `p`.
+Print progress bar to `io`.
 
 # Arguments
 - `io::IO`
-- `p::ProgressBar`
+- `barglyphs::BarGlyphs`
+- `tfirst::Float64`
 - `desc`: description to be printed at left side of progress bar.
 - `progress::Real`: a number between 0 and 1 or a `NaN`.
+- `eta_seconds::Real`: ETA in seconds
 """
-function printprogress(io::IO, p::ProgressBar, desc, progress::Real)
+function printprogress(
+    io::IO,
+    barglyphs::BarGlyphs,
+    tfirst::Float64,
+    desc,
+    progress::Real,
+    eta_seconds::Real,
+)
     t = time()
     percentage_complete = 100.0 * (isnan(progress) ? 0.0 : progress)
 
     #...length of percentage and ETA string with days is 29 characters
     barlen = max(0, displaysize(io)[2] - (length(desc) + 29))
 
     if progress >= 1
-        bar = barstring(barlen, percentage_complete, barglyphs=p.barglyphs)
-        dur = durationstring(t - p.tfirst)
+        bar = barstring(barlen, percentage_complete, barglyphs=barglyphs)
+        dur = durationstring(t - tfirst)
         @printf io "%s%3u%%%s Time: %s" desc round(Int, percentage_complete) bar dur
         return
     end
 
-    bar = barstring(barlen, percentage_complete, barglyphs=p.barglyphs)
-    elapsed_time = t - p.tfirst
-    est_total_time = 100 * elapsed_time / percentage_complete
-    if 0 <= est_total_time <= typemax(Int)
-        eta_sec = round(Int, est_total_time - elapsed_time)
+    bar = barstring(barlen, percentage_complete, barglyphs=barglyphs)
+    if 0 <= eta_seconds <= typemax(Int)
+        eta_sec = round(Int, eta_seconds)
         eta = durationstring(eta_sec)
     else
         eta = "N/A"

src/StickyMessages.jl

@@ -95,7 +95,7 @@ end
 
 function Base.push!(sticky::StickyMessages, message::Pair)
     if !sticky.ansi_codes
-        write(sticky.io, message[2])
+        println(sticky.io, rstrip(message[2]))
         return
     end
     label,text = message

src/TerminalLogger.jl

@@ -29,7 +29,7 @@ struct TerminalLogger <: AbstractLogger
     right_justify::Int
     message_limits::Dict{Any,Int}
     sticky_messages::StickyMessages
-    bars::Dict{Any,ProgressBar}
+    bartrees::Vector{Node{ProgressBar}}
 end
 function TerminalLogger(stream::IO=stderr, min_level=ProgressLevel;
                         meta_formatter=default_metafmt, show_limited=true,
@@ -42,7 +42,7 @@ function TerminalLogger(stream::IO=stderr, min_level=ProgressLevel;
         right_justify,
         Dict{Any,Int}(),
         StickyMessages(stream),
-        Dict{Any,ProgressBar}(),
+        Union{}[],
     )
 end
 
@@ -104,61 +104,105 @@ function termlength(str)
     return N
 end
 
-function handle_progress(logger, message, id, progress)
-    # Don't do anything when it's already done:
-    if (progress == "done" || progress >= 1) && !haskey(logger.bars, id)
-        return
+function findbar(bartree, id)
+    if !(bartree isa AbstractArray)
+        bartree.data.id === id && return bartree
+    end
+    for node in bartree
+        found = findbar(node, id)
+        found === nothing || return found
     end
+    return nothing
+end
 
-    try
-        bar = get!(ProgressBar, logger.bars, id)
+function foldtree(op, acc, tree)
+    for node in tree
+        acc = foldtree(op, acc, node)
+    end
+    if !(tree isa AbstractArray)
+        acc = op(acc, tree)
+    end
+    return acc
+end
 
-        if message == ""
-            message = "Progress: "
-        else
-            message = string(message)
-            if !endswith(message, " ")
-                message *= " "
-            end
-        end
+const BAR_MESSAGE_ID = gensym(:BAR_MESSAGE_ID)
 
-        bartxt = sprint(
-            printprogress,
-            bar,
-            message,
-            progress == "done" ? 1.0 : progress;
-            context = :displaysize => displaysize(logger.stream),
-        )
+function handle_progress(logger, progress)
+    node = findbar(logger.bartrees, progress.id)
+    if node === nothing
+        # Don't do anything when it's already done:
+        (progress.done || something(progress.fraction, 0.0) >= 1) && return
 
-        if progress == "done" || progress >= 1
-            pop!(logger.sticky_messages, id)
-            printstyled(logger.stream, bartxt; color=:light_black)
-            println(logger.stream)
+        parentnode = findbar(logger.bartrees, progress.parentid)
+        bar = ProgressBar(
+            fraction = progress.fraction,
+            name = progress.name,
+            id = progress.id,
+            parentid = progress.parentid,
+        )
+        if parentnode === nothing
+            node = Node(bar)
+            pushfirst!(logger.bartrees, node)
         else
-            bartxt = sprint(context = logger.stream) do io
-                printstyled(io, bartxt; color=:green)
-            end
-            push!(logger.sticky_messages, id => bartxt)
+            bar.level = parentnode.data.level + 1
+            node = addchild(parentnode, bar)
         end
-    finally
-        if progress == "done" || progress >= 1
-            pop!(logger.sticky_messages, id)  # redundant (but safe) if no error
-            pop!(logger.bars, id, nothing)
+    else
+        bar = node.data
+        set_fraction!(bar, progress.fraction)
+        if progress.name != ""
+            bar.name = progress.name
+        end
+        node.data = bar
+    end
+    if progress.done
+        if isroot(node)
+            deleteat!(logger.bartrees, findfirst(x -> x === node, logger.bartrees))
+        else
+            prunebranch!(node)
+        end
+    end
+
+    bartxt = sprint(context = :displaysize => displaysize(logger.stream)) do io
+        foldtree(true, logger.bartrees) do isfirst, node
+            isfirst || println(io)
+            printprogress(io, node.data)
+            false  # next `isfirst`
+        end
+    end
+    if isempty(bartxt)
+        pop!(logger.sticky_messages, BAR_MESSAGE_ID)
+    else
+        bartxt = sprint(context = logger.stream) do io
+            printstyled(io, bartxt; color=:green)
+        end
+        push!(logger.sticky_messages, BAR_MESSAGE_ID => bartxt)
+    end
+
+    # "Flushing" non-sticky message should be done after the sticky
+    # message is re-drawn:
+    if progress.done
+        ensure_done!(bar)
+        donetxt = sprint(context = :displaysize => displaysize(logger.stream)) do io
+            printprogress(io, bar)
         end
+        printstyled(logger.stream, donetxt; color=:light_black)
+        println(logger.stream)
     end
 end
 
 function handle_message(logger::TerminalLogger, level, message, _module, group, id,
-                        filepath, line; maxlog=nothing, progress=nothing,
+                        filepath, line; maxlog=nothing,
                         sticky=nothing, kwargs...)
     if maxlog !== nothing && maxlog isa Integer
         remaining = get!(logger.message_limits, id, maxlog)
         logger.message_limits[id] = remaining - 1
         remaining > 0 || return
     end
 
-    if progress == "done" || progress isa Real
-        handle_progress(logger, message, id, progress)
+    progress = asprogress(level, message, _module, group, id, filepath, line; kwargs...)
+    if progress !== nothing
+        handle_progress(logger, progress)
         return
     end
 

src/TerminalLoggers.jl

@@ -7,15 +7,20 @@ using Logging:
 import Logging:
     handle_message, shouldlog, min_enabled_level, catch_exceptions
 
+using LeftChildRightSiblingTrees: Node, addchild, isroot, prunebranch!
+using ProgressLogging: asprogress
+using UUIDs: UUID
+
 export TerminalLogger
 
 const ProgressLevel = LogLevel(-1)
 
 include("ProgressMeter/ProgressMeter.jl")
 using .ProgressMeter:
-    ProgressBar, printprogress
+    BarGlyphs
 
 include("StickyMessages.jl")
+include("ProgressBar.jl")
 include("TerminalLogger.jl")
 
 end # module