Support general I/O for ProgressBar output

src/progress_bar.jl

@@ -1,4 +1,4 @@
-export ProgressBar
+export ProgressBar, next!
 
 struct ProgressBar{CT,T1<:Integer, T2<:Real}
     counter::CT
@@ -9,11 +9,11 @@ struct ProgressBar{CT,T1<:Integer, T2<:Real}
     lock::ReentrantLock
 end
 
-function ProgressBar(max_counts; enable=true, bar_width=30)
+function ProgressBar(max_counts::Int; enable::Bool=true, bar_width::Int=30)
     return ProgressBar(Ref{Int64}(0), max_counts, enable, bar_width, time(), ReentrantLock())
 end
 
-function next!(p::ProgressBar)
+function next!(p::ProgressBar, io::IO=stdout)
 
     lock(p.lock)
 
@@ -43,10 +43,10 @@ function next!(p::ProgressBar)
     # Construct the progress bar string
     bar = "[" * repeat("=", progress) * repeat(" ", bar_width - progress) * "]"
 
-    print("\rProgress: $bar $percentage_100% --- Elapsed Time: $elapsed_time_str (ETA: $eta_str)")
-    flush(stdout)
+    print(io, "\rProgress: $bar $percentage_100% --- Elapsed Time: $elapsed_time_str (ETA: $eta_str)")
+    flush(io)
 
     unlock(p.lock)
 
-    p.counter[] >= p.max_counts ? print("\n") : nothing
+    p.counter[] >= p.max_counts ? print(io, "\n") : nothing
 end

test/low_rank_dynamics.jl

@@ -66,7 +66,9 @@
         p0         = 0.,
         atol_inv   = 1e-6,
         adj_condition="variational",
-        Δt = 0.2, )
+        Δt = 0.2, 
+        progress = false
+    )
     lrsol = lr_mesolve(H, z, B, tl, c_ops; e_ops=e_ops, f_ops=(f_entropy,), opt=opt)
 
     # Test

test/progress_bar.jl

@@ -0,0 +1,14 @@
+@testset "Progress Bar" begin
+    bar_width = 30
+    strLength = 67 + bar_width # including "\r" in the beginning of the string
+    prog = ProgressBar(bar_width, enable=true, bar_width=bar_width)
+    for p in 1:bar_width
+        output = sprint((t, s) -> next!(s, t), prog)
+
+        if p < bar_width
+            @test length(output) == strLength
+        else # the last output has an extra "\n" in the end
+            @test length(output) == strLength + 1
+        end
+    end
+end

test/runtests.jl

@@ -5,8 +5,8 @@ const GROUP = get(ENV, "GROUP", "All")
 
 const testdir = dirname(@__FILE__)
 
-# Put them in alphabetical order
-tests = [
+# Put core tests in alphabetical order
+core_tests = [
     "correlations_and_spectrum.jl",
     "dynamical_fock_dimension_mesolve.jl",
     "dynamical-shifted-fock.jl",
@@ -16,14 +16,15 @@ tests = [
     "low_rank_dynamics.jl",
     "negativity_and_partial_transpose.jl",
     "permutation.jl",
+    "progress_bar.jl",
     "quantum_objects.jl",
     "steady_state.jl",
     "time_evolution_and_partial_trace.jl",
     "wigner.jl",
 ]
 
 if (GROUP == "All") || (GROUP == "Core")
-    for test in tests
+    for test in core_tests
         include(joinpath(testdir, test))
     end
 end