From 79bdd7a4ace644aada0953f2cc136ddfe440ef4c Mon Sep 17 00:00:00 2001 From: Paul Berg Date: Mon, 11 Jan 2021 19:26:35 +0100 Subject: [PATCH] handle filters in generators --- src/analysis/ExpressionExplorer.jl | 3 +++ test/ExpressionExplorer.jl | 3 +++ 2 files changed, 6 insertions(+) diff --git a/src/analysis/ExpressionExplorer.jl b/src/analysis/ExpressionExplorer.jl index 468ad48772..a281ef142f 100644 --- a/src/analysis/ExpressionExplorer.jl +++ b/src/analysis/ExpressionExplorer.jl @@ -357,6 +357,9 @@ function explore!(ex::Expr, scopestate::ScopeState)::SymbolsState elseif ex.head == :let || ex.head == :for || ex.head == :while # Creates local scope return explore_inner_scoped(ex, scopestate) + elseif ex.head == :filter + # In a filter, the assignment is the second expression, the condition the first + return mapfoldr(a -> explore!(a, scopestate), union!, ex.args, init=SymbolsState()) elseif ex.head == :generator # Creates local scope diff --git a/test/ExpressionExplorer.jl b/test/ExpressionExplorer.jl index 9568722ea7..409ac02814 100644 --- a/test/ExpressionExplorer.jl +++ b/test/ExpressionExplorer.jl @@ -119,6 +119,9 @@ using Test @test testee(:([sqrt(s) for s in 1:n]), [:n], [], [:sqrt, :(:)], []) @test testee(:([sqrt(s + r) for s in 1:n, r in k]), [:n, :k], [], [:sqrt, :(:), :+], []) @test testee(:([s + j + r + m for s in 1:3 for j in 4:5 for (r, l) in [(1, 2)]]), [:m], [], [:+, :(:)], []) + @test testee(:([a for a in b if a != 2]), [:b], [], [:(!=)], []) + @test testee(:([a for a in f() if g(a)]), [], [], [:f, :g], []) + @test testee(:([c(a) for a in f() if g(a)]), [], [], [:c, :f, :g], []) @test testee(:([a for a in a]), [:a], [], [], []) @test testee(:(for a in a; a; end), [:a], [], [], [])