Preserve offset axes in @progress with comprehension

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))) = _vec_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 = first.(only.(axes.(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
 
 
@@ -138,3 +139,6 @@ end
 
 _comprehension(iter_exprs, body,) = Expr(:comprehension, Expr(:generator, body, iter_exprs...))
 _for(iter_exprs, body) = Expr(:for, Expr(:block, reverse(iter_exprs)...), body)
+
+# like Base.pairs, but for any vector with `axes`
+_vec_pairs(v) = zip(only(axes(v)), v)

test/runtests.jl

@@ -1,5 +1,6 @@
 using Juno
 using Test
+using OffsetArrays
 
 @test Juno.isactive() == false
 
@@ -64,4 +65,31 @@ 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
+
+# non-indexable iterables with axes
+let r = 1:5
+  x1 = @progress y1 = [i for i in (x^2 for x in r)]
+  x2 = @progress y2 = [i for i in zip(r,r)]
+  @test x1 == y1 == r.^2
+  @test x2 == y2 == collect(zip(r,r))
+
+  y1, y2 = [], []
+  x1 = @progress for i in (x^2 for x in r)
+    push!(y1, i)
+  end
+  x2 = @progress for i in zip(r,r)
+    push!(y2, i)
+  end
+  @test x1 == x2 == nothing
+  @test y1 == r.^2
+  @test y2 == collect(zip(r,r))
+end
+
+
 @test Juno.notify("hi") == nothing