Skip to content

Instantly share code, notes, and snippets.

@Ismael-VC
Last active January 14, 2020 19:24
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save Ismael-VC/e3f74338b2c9ba34cec7fe3c46be05be to your computer and use it in GitHub Desktop.
Save Ismael-VC/e3f74338b2c9ba34cec7fe3c46be05be to your computer and use it in GitHub Desktop.
`vcomp` macro: vector comprehensions with conditional `when` clause for Julia.
"""
@vcomp(comprehension::Expr, when::Symbol, condition)
`@vcomp` (*vector comprehension*) based on Python's list comprehensions
with a conditional clause to filter elements in a succinct way.
# Examples
```julia
julia> @vcomp Int[i^2 for i in 1:10] when i % 2 == 0
5-element Array{Int,1}:
4
16
36
64
100
julia> words = ["ant", "akita", "bear", "cow"];
julia> @vcomp [uppercase(word) for word in words] when startswith(word, "a")
2-element Array{Any,1}:
"ANT"
"AKITA"
julia> @vcomp Tuple[(i, e) for (i, e) in enumerate(reverse(1:10))] when i-e < 0
5-element Array{Tuple,1}:
(1,10)
(2,9)
(3,8)
(4,7)
(5,6)
```
"""
macro vcomp(comprehension::Expr, when::Symbol, condition)
comp_head, comp_args = comprehension.head, comprehension.args
comp_head ∉ [:comprehension, :typed_comprehension] && error("@vcomp not a comprehension")
when ≠ :when && error("@vcomp expected `when`, got: `$when`")
T = comp_head == :typed_comprehension ? comp_args[1] : nothing
if VERSION < v"0.5-"
element = comp_head == :comprehension ? comp_args[1] : comp_args[2]
sequence = comp_head == :comprehension ? comp_args[2] : comp_args[3]
else
element = comp_head == :comprehension ? comp_args[1].args[1] : comp_args[2].args[1]
sequence = comp_head == :comprehension ? comp_args[1].args[2] : comp_args[2].args[2]
end
result = T ≠ nothing ? :($T[]) : :([])
block = Expr(:let, Expr(:block,
Expr(:(=), :res, result),
Expr(:for, sequence,
Expr(:if, condition,
Expr(:call, :push!, :res, element))),
:res))
return esc(block)
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment