Preserve offset axes in @progress with comprehension

Author: yha

Project.toml

@@ -13,7 +13,8 @@ Media = "0.5"
 julia = "0.7, 1"
 
 [extras]
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Test", "OffsetArrays"]

src/progress.jl

@@ -98,7 +98,7 @@ end
 
 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))))
+  iter_exprs = [:(($i,$(esc(v))) = pairs($(esc(r))))
                   for (i,v,r) in zip(count_vars,iter_vars,ranges)]
   _id = "progress_$(gensym())"
   quote
@@ -107,10 +107,11 @@ function _progress(name, thresh, ex, target, result, loop, iter_vars, ranges, bo
       $target = try
         ranges = $(Expr(:vect,esc.(ranges)...))
         nranges = length(ranges)
+        starts = firstindex.(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
+        _frac(i) = (sum((i-j)*s for (i,j,s) in zip(i,starts,strides)) + 1) / n
         lastfrac = 0.0
 
 

test/runtests.jl

@@ -1,5 +1,6 @@
 using Juno
 using Test
+using OffsetArrays
 
 @test Juno.isactive() == false
 
@@ -64,4 +65,11 @@ let a = [], x
   @test x == nothing
 end
 
+let off1 = -2, off2 = 21
+  v1 = OffsetArray(1:3, off1)
+  v2 = OffsetArray(-1:1, off2)
+  x = @progress y = [i*j for i in v1, j in v2]
+  @test x == y == OffsetArray([-1 0 1; -2 0 2; -3 0 3], off1, off2)
+end
+
 @test Juno.notify("hi") == nothing