Skip to content

Commit 7e11769

Browse files
fix: add support for comprehensions with conditionals
Comprehension support works by parsing the underlying generators, and generators with `if` conditions have `filter` as their `Expr` form's second argument. This previously wasn't supported, as the generator parsing code expected the `for` iteration's `=` to be the second arg instead. This change adds support for the case of generators with filters, and adds relevant tests.
1 parent df601ea commit 7e11769

File tree

3 files changed

+40
-8
lines changed

3 files changed

+40
-8
lines changed

CHANGELOG.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
# News
22

3+
## v1.0.4 - 2025-08-16
4+
5+
- Support comprehensions with conditionals, e.g., `@resumable function f1(); [i for i in 1:10 if i<5]; end` which previously led to "Illegal expression" errors.
6+
37
## v1.0.3 - 2025-03-24
48

59
- Internal changes to `fsmi_generator` to support julia 1.12

src/utils.jl

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -338,19 +338,31 @@ function scoping(s::Symbol, scope; new = false)
338338
return lookup_rhs!(s, scope)
339339
end
340340

341-
function scope_generator(expr, scope)
342-
expr.head !== :generator && error("Illegal generator expression: $(expr)")
343-
# first the generator case
341+
function scope_generator_inner(expr, scope)
344342
for i in 2:length(expr.args)
345-
!(expr.args[i] isa Expr && expr.args[i].head === :(=)) && error("Illegal expression in generator: $(expr.args[i])")
343+
!(expr.args[i] isa Expr && expr.args[i].head === :(=)) &&
344+
error("Illegal expression in generator: $(expr.args[i])")
346345
expr.args[i].args[2] = scoping(expr.args[i].args[2], scope)
347346
end
348347
# now create new scope
349348
push!(scope.scope_stack, Dict())
350349
for i in 2:length(expr.args)
351350
expr.args[i].args[1] = lookup_lhs!(expr.args[i].args[1], scope, new = true)
352351
end
352+
end
353+
354+
function scope_generator(expr, scope)
355+
expr.head !== :generator && error("Illegal generator expression: $(expr)")
353356

357+
has_filter = length(expr.args) == 2 && expr.args[2] isa Expr && expr.args[2].head === :filter
358+
if has_filter
359+
ex = expr.args[2]
360+
scope_generator_inner(ex, scope)
361+
# now apply scoping to the filter condition expression
362+
ex.args[1] = scoping(ex.args[1], scope)
363+
else
364+
scope_generator_inner(expr, scope)
365+
end
354366
expr.args[1] = scoping(expr.args[1], scope)
355367
pop!(scope.scope_stack)
356368
return expr
@@ -432,10 +444,6 @@ function scoping(expr::Expr, scope)
432444
end
433445
return expr
434446
end
435-
if expr.head === :generator
436-
expr = scope_generator(expr)
437-
return expr
438-
end
439447

440448
if expr.head === :(=)
441449
# One special case, where we need to have both LHS and RHS at our hands

test/test_main.jl

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -403,6 +403,26 @@ end
403403
end
404404
end
405405
@test collect(test_comprehension5()) == [2, 3, 4, 3, 4, 5]
406+
407+
@resumable function test_comprehension_if()
408+
@yield [i for i in 1:10 if i < 5]
409+
end
410+
@test collect(test_comprehension_if()) == [[1, 2, 3, 4]]
411+
412+
@resumable function test_comprehension_if_multi_iterators()
413+
@yield [i + j for i in 1:3, j in 1:3 if i <= j]
414+
end
415+
@test collect(test_comprehension_if_multi_iterators()) == [[2, 3, 4, 4, 5, 6]]
416+
417+
@resumable function test_comprehension_if_nested_for()
418+
@yield [i * j for i in 1:3 for j in 1:3 if i * j < 5]
419+
end
420+
@test collect(test_comprehension_if_nested_for()) == [[1, 2, 3, 2, 4, 3]]
421+
422+
@resumable function test_comprehension_if_mixed()
423+
@yield [i + j + k for i in 1:2, j in 1:2 for k in 1:2 if i + j + k == 4]
424+
end
425+
@test collect(test_comprehension_if_mixed()) == [[4, 4, 4]]
406426
end
407427

408428
@testset "test_ref" begin

0 commit comments

Comments
 (0)