Preserve offset axes in @progress with comprehension

Project.toml

@@ -11,8 +11,10 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 julia = "0.7, 1"
 
 [extras]
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+
 [targets]
-test = ["Test", "Random"]
+test = ["OffsetArrays", "Random", "Test"]

src/ProgressLogging.jl

@@ -452,7 +452,7 @@ end
 
 function _progress(name, thresh, ex, target, result, loop, iter_vars, ranges, body)
     count_vars = [gensym(Symbol("i$k")) for k = 1:length(iter_vars)]
-    iter_exprs = [:(($i, $v) = $enumerate($r)) for (i, v, r) in zip(
+    iter_exprs = [:(($i, $v) = $zip($_linindex($r),$r)) for (i, v, r) in zip(
         count_vars,
         iter_vars,
         ranges,
@@ -496,14 +496,17 @@ end
 
 function make_count_to_frac(iterators...)
     lens = map(length, iterators)
+    firsts = map(first∘_linindex, iterators)
     n = prod(lens)
     strides = (1, taccumulate(*, Base.front(lens))...)
     function count_to_frac(idxs...)
-        offsets = map(i -> i - 1, idxs)
+        offsets = map(-, idxs, firsts)
         total = sum(map(*, offsets, strides)) + 1
         return total / n
     end
     return count_to_frac
 end
 
+_linindex(a) = LinearIndices(axes(a))
+
 end # module

test/test_progress_macro.jl

@@ -1,63 +1,105 @@
 module TestProgressMacro
 
-using ProgressLogging: @progress
+using ProgressLogging: @progress, ProgressLevel
 using Test
+using Test: collect_test_logs
+using OffsetArrays
 
-let i = 0, x
-    x = @progress for _ = 1:100
-        i += 1
-    end
-    @test i == 100
-    @test x == nothing
-end
+@testset "@progress" begin
+  let i = 0, x
+      x = @progress for _ = 1:100
+          i += 1
+      end
+      @test i == 100
+      @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, 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
-    end
-    @test i == 100
-    @test x == nothing
-end
+  let i = 0, x
+      x = @progress "named" for _ = 1:100
+          i += 1
+      end
+      @test i == 100
+      @test x == nothing
+  end
 
-let i = 0, j = 0, x
-    x = @progress for _ = 1:10, __ = 1:20
-        i += 1
-    end
-    @test i == 200
-    @test x == nothing
-end
+  let i = 0, j = 0, x
+      x = @progress for _ = 1:10, __ = 1:20
+          i += 1
+      end
+      @test i == 200
+      @test x == nothing
+  end
+
+  let i = 0, j = 0, x
+      bar = "bar"
+      x = @progress "foo $bar" for _ = 1:10
+          i += 1
+      end
+      @test i == 10
+      @test x == nothing
+  end
+
+  let x, y
+      x = @progress y = [i + 3j for i = 1:3, j = 1:4]
+      @test y == reshape(4:15, 3, 4)
+      @test x == y
+  end
 
-let i = 0, j = 0, x
-    bar = "bar"
-    x = @progress "foo $bar" for _ = 1:10
-        i += 1
+  let a = [], x
+      x = @progress for i = 1:3, j in [-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
+
+  # Multi-dimensional arrays in comprehension and offset axes
+  let off1 = -2, off2 = 21
+    v1 = OffsetArray(2:3, off1)
+    v2 = OffsetArray(-1:2, off2)
+    logs, _ = collect_test_logs(min_level = ProgressLevel) do
+      x = @progress y = [i*j for i in v1, j in v2]
+      @test x == y == OffsetArray([-2 0 2 4; -3 0 3 6], off1, off2)
     end
-    @test i == 10
-    @test x == nothing
-end
+    @test isequal(
+        [l.kwargs[:progress] for l in logs],
+        [nothing; (1:8)./8; "done"],
+    )
+    m = OffsetArray(reshape(2:7,2,3), off1, off2)
+    x = @progress y = [i*j for i in v1, j in m]
+    @test x == y == [i*j for i in v1, j in m]
+  end
 
-let x, y
-    x = @progress y = [i + 3j for i = 1:3, j = 1:4]
-    @test y == reshape(4:15, 3, 4)
-    @test x == y
-end
+  # non-indexable iterables with axes
+  @testset "non-indexable" for off in (0,10)
+    let r = OffsetVector(1:5, off)
+      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))
 
-let a = [], x
-    x = @progress for i = 1:3, j in [-5, -2, -1, 8]
-        j > 0 && continue
-        push!(a, (i, j))
-        i > 1 && break
+      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 OffsetVector(y1,off) == r.^2
+      @test OffsetVector(y2,off) == collect(zip(r,r))
     end
-    @test a == [(1, -5), (1, -2), (1, -1), (2, -5)]
-    @test x == nothing
+  end
 end
 
 end  # module