@progress fixes (#166)

* Progress bar: multiple iteration variables & comprehensions * Progress bar: multiple iteration variables & comprehensions * Fix argument to _frac * More progress bar tests * More detailed error message. Simplified by inlining _update. * Fix global scope assignment * Escape ranges * Test non-literal range * for loops in @progress not rewritten as comprehension, to support break, continue
JunoLab · Sep 4, 2018 · fa9a3ad · fa9a3ad
1 parent 3db0633
commit fa9a3ad
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 16 deletions.
diff --git a/src/progress.jl b/src/progress.jl
@@ -59,6 +59,7 @@ function _progress(name, thresh, ex)
         ex.args[2].head == :comprehension &&
         ex.args[2].args[1].head == :generator
     # comprehension: <target> = [<body> for <iter_var> in <range>,...]
+    loop = _comprehension
     target = esc(ex.args[1])
     result = target
     gen_ex = ex.args[2].args[1]
@@ -70,6 +71,7 @@ function _progress(name, thresh, ex)
         ex.args[1].head == Symbol("=") &&
         ex.args[2].head == :block
     # single-variable for: for <iter_var> = <range>; <body> end
+    loop = _for
     target = :_
     result = :nothing
     iter_vars = [ex.args[1].args[1]]
@@ -79,19 +81,22 @@ function _progress(name, thresh, ex)
         ex.args[1].head == :block &&
         ex.args[2].head == :block
     # multi-variable for: for <iter_var> = <range>,...; <body> end
+    loop = _for
     target = :_
     result = :nothing
-    iter_vars = [e.args[1] for e in ex.args[1].args]
-    ranges = [e.args[2] for e in ex.args[1].args]
+    # iter_vars and ranges are ordered from inner loop to outer loop, for
+    # consistent computation of progress between for loops and comprehensions
+    iter_vars = reverse([e.args[1] for e in ex.args[1].args])
+    ranges = reverse([e.args[2] for e in ex.args[1].args])
     body = esc(ex.args[2])
   else
     error("@progress requires a for loop (for i in irange, j in jrange, ...; <body> end) " *
           "or array comprehension with assignment (x = [<body> for i in irange, j in jrange, ...])")
   end
-  _progress(name, thresh, ex, target, result, iter_vars, ranges, body)
+  _progress(name, thresh, ex, target, result, loop, iter_vars, ranges, body)
 end
 
-function _progress(name, thresh, ex, target, result, iter_vars, ranges, body)
+function _progress(name, thresh, ex, target, result, loop, iter_vars, ranges, body)
   count_vars = [Symbol("i$k") for k=1:length(iter_vars)]
   iter_exprs = [:(($i,$(esc(v))) = enumerate($(esc(r))))
                   for (i,v,r) in zip(count_vars,iter_vars,ranges)]
@@ -100,25 +105,26 @@ function _progress(name, thresh, ex, target, result, iter_vars, ranges, body)
     if isactive()
       @logmsg($PROGRESSLEVEL, $name, progress=0.0, _id=Symbol($_id))
       $target = try
-        ranges = $(Expr(:vect,ranges...))
+        ranges = $(Expr(:vect,esc.(ranges)...))
         nranges = length(ranges)
         lens = length.(ranges)
         n = prod(lens)
         strides = cumprod([1;lens[1:end-1]])
         _frac(i) = (sum((i-1)*s for (i,s) in zip(i,strides)) + 1) / n
         lastfrac = 0.0
 
-        $(Expr(:comprehension, Expr(:generator,
-                            quote
-                              frac = _frac($(Expr(:vect, count_vars...)))
-                              if frac - lastfrac > $thresh
-                                @logmsg($PROGRESSLEVEL, $name, progress=frac, _id=Symbol($_id))
-                                lastfrac = frac
-                              end
-                              $body
-                            end,
-                            iter_exprs...
-                         )))
+
+        $(loop(iter_exprs,
+            quote
+                frac = _frac($(Expr(:vect, count_vars...)))
+                if frac - lastfrac > $thresh
+                    @logmsg($PROGRESSLEVEL, $name, progress=frac, _id=Symbol($_id))
+                    lastfrac = frac
+                end
+                $body
+            end
+        ))
+
       finally
         @logmsg($PROGRESSLEVEL, $name, progress="done", _id=Symbol($_id))
       end
@@ -128,3 +134,6 @@ function _progress(name, thresh, ex, target, result, iter_vars, ranges, body)
     end
   end
 end
+
+_comprehension(iter_exprs, body,) = Expr(:comprehension, Expr(:generator, body, iter_exprs...))
+_for(iter_exprs, body) = Expr(:for, Expr(:block, reverse(iter_exprs)...), body)
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -15,6 +15,14 @@ let i = 0, x
   @test x == nothing
 end
 
+let i = 0, r = -50:10:50, x
+  x = @progress for _ in r
+    i += 1
+  end
+  @test i == 11
+  @test x == nothing
+end
+
 let i = 0, x
   x = @progress "named" for _ = 1:100
     i += 1
@@ -28,6 +36,7 @@ let i = 0, j = 0, x
     i += 1
   end
   @test i == 200
+  @test x == nothing
 end
 
 let x,y
@@ -36,5 +45,14 @@ let x,y
   @test x == y
 end
 
+let a = [], x
+  x = @progress for i=1:3, j=[-5,-2,-1,8]
+      j > 0 && continue
+      push!(a,(i,j))
+      i > 1 && break
+  end
+  @test a == [(1,-5),(1,-2),(1,-1),(2,-5)]
+  @test x == nothing
+end
 
 @test Juno.notify("hi") == nothing