Skip to content

Commit

Permalink
Fix for non-array iterables
Browse files Browse the repository at this point in the history
  • Loading branch information
yha committed Sep 1, 2020
1 parent bc5a451 commit d982fb7
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 2 deletions.
7 changes: 5 additions & 2 deletions src/progress.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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))) = pairs($(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
Expand All @@ -107,7 +107,7 @@ 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)
starts = first.(only.(axes.(ranges)))
lens = length.(ranges)
n = prod(lens)
strides = cumprod([1;lens[1:end-1]])
Expand Down Expand Up @@ -139,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)
20 changes: 20 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -72,4 +72,24 @@ let off1 = -2, off2 = 21
@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

0 comments on commit d982fb7

Please sign in to comment.