Skip to content

Commit dce0526

Browse files
committed
added short form syntax, fixes #6
1 parent ba713e3 commit dce0526

File tree

2 files changed

+92
-33
lines changed

2 files changed

+92
-33
lines changed

src/KeywordDispatch.jl

Lines changed: 73 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,22 @@ argtype(x::Symbol) = Any
5151
argtype(x::Expr) = x.head == :(::) ? x.args[end] : error("unexpected expression $x")
5252

5353

54+
function unwrap_where(expr)
55+
stack = Any[]
56+
while expr isa Expr && expr.head == :where
57+
push!(stack, expr.args[2])
58+
expr = expr.args[1]
59+
end
60+
expr, stack
61+
end
62+
63+
function wrap_where(expr, stack)
64+
for w in Iterators.reverse(stack)
65+
expr = Expr(:where, expr, esc(w))
66+
end
67+
expr
68+
end
69+
5470
struct KeywordMethodError <: Exception
5571
f
5672
args
@@ -88,14 +104,16 @@ end
88104

89105

90106
"""
91-
@kwdispatch expr
107+
@kwdispatch sig [methods]
92108
93-
Designate a function signature `expr` that should dispatch on keyword arguments. A
109+
Designate a function signature `sig` that should dispatch on keyword arguments. A
94110
function can also be provided, in which case _all_ calls to that function are will
95111
dispatch to keyword methods.
96112
97-
Note that no keywords should appear in `@kwdispatch` signatures. To define the keyword
98-
methods, use the [`@kwmethod`](@ref) macro.
113+
Note that no keywords should appear in `sig` signatures.
114+
115+
The optional `methods` argument allows a block of keyword methods specified as anonymous
116+
functions. To define additional keyword methods, use the [`@kwmethod`](@ref) macro.
99117
100118
# Examples
101119
@@ -105,17 +123,18 @@ methods, use the [`@kwmethod`](@ref) macro.
105123
106124
@kwdispacth f # equivalent to @kwdispatch f(_...)
107125
126+
@kwdispatch f(x) begin
127+
(a) -> x+a
128+
(b) -> x-b
129+
end
130+
# equivalent to
131+
# @kwdispatch f(x)
132+
# @kwmethod f(x;a) = x+a
133+
# @kwmethod f(x;b) = x-b
108134
```
109135
"""
110-
macro kwdispatch(fexpr)
111-
fexpr = outexpr = :($fexpr = _)
112-
113-
# unwrap where clauses
114-
while fexpr.args[1] isa Expr && fexpr.args[1].head == :where
115-
fexpr = fexpr.args[1]
116-
fexpr.args[2] = esc(fexpr.args[2])
117-
end
118-
fcall = fexpr.args[1]
136+
macro kwdispatch(fexpr,methods=nothing)
137+
fcall, wherestack = unwrap_where(fexpr)
119138

120139
# handle: `fun`, `Mod.fun`, `(a::B)`
121140
if fcall isa Symbol || fcall isa Expr && (fcall.head in (:., :(::), :curly))
@@ -125,8 +144,8 @@ macro kwdispatch(fexpr)
125144
@assert fcall isa Expr && fcall.head == :call
126145

127146
f = fcall.args[1]
128-
fargs = fcall.args[2:end]
129-
if length(fargs) >= 1 && fargs[1] isa Expr && fargs[1].head == :parameters
147+
posargs = fcall.args[2:end]
148+
if length(posargs) >= 1 && posargs[1] isa Expr && posargs[1].head == :parameters
130149
error("keyword arguments should only appear in @kwdispatch expressions")
131150
end
132151
f = argmeth(f)
@@ -137,13 +156,35 @@ macro kwdispatch(fexpr)
137156
ftype = :(typeof($(esc(f))))
138157
end
139158

140-
fargs_method = argmeth.(fargs)
141-
fexpr.args[1] = :($(esc(f))($(esc.(fargs_method)...); kwargs...))
159+
posargs_method = argmeth.(posargs)
160+
ff = esc(argsym(f))
142161

143-
outexpr.args[2] = :(kwcall(ntsort(kwargs.data), $(esc(argsym(f))), $(esc.(argsym.(fargs_method))...)))
144-
return outexpr
162+
quote
163+
$(wrap_where(:($(esc(f))($(esc.(posargs_method)...); kwargs...)), wherestack)) =
164+
KeywordDispatch.kwcall(ntsort(kwargs.data), $ff, $(esc.(argsym.(posargs_method))...))
165+
$(generate_kwmethods(methods, f, posargs, wherestack))
166+
end
145167
end
146168

169+
generate_kwmethods(other, f, posargs, wherestack) = other
170+
function generate_kwmethods(expr::Expr, f, posargs, wherestack)
171+
if expr.head == :block
172+
for (i, ex) in enumerate(expr.args)
173+
expr.args[i] = generate_kwmethods(ex, f, posargs, wherestack)
174+
end
175+
return expr
176+
elseif expr.head in (:->, :function)
177+
if expr.args[1] isa Symbol
178+
return kwmethod_expr(f, posargs, [expr.args[1]], wherestack, esc(expr.args[2]))
179+
elseif expr.args[1] isa Expr && expr.args[1].head == :tuple
180+
return kwmethod_expr(f, posargs, expr.args[1].args, wherestack, esc(expr.args[2]))
181+
end
182+
end
183+
error("Invalid keyword definition $expr.")
184+
end
185+
186+
187+
147188
"""
148189
@kwmethod expr
149190
@@ -165,16 +206,10 @@ The positional signature should first be designated by the [`@kwdispatch`](@ref)
165206
"""
166207
macro kwmethod(fexpr)
167208
@assert fexpr isa Expr && fexpr.head in (:function, :(=))
168-
fexpr.args[2] = esc(fexpr.args[2])
209+
body = esc(fexpr.args[2])
169210

170-
outexpr = fexpr
171-
# unwrap where clauses
172-
while fexpr.args[1] isa Expr && fexpr.args[1].head == :where
173-
fexpr = fexpr.args[1]
174-
fexpr.args[2] = esc(fexpr.args[2])
175-
end
211+
fcall, wherestack = unwrap_where(fexpr.args[1])
176212

177-
fcall = fexpr.args[1]
178213
@assert fcall isa Expr && fcall.head == :call
179214

180215
f = fcall.args[1]
@@ -183,8 +218,12 @@ macro kwmethod(fexpr)
183218
error("@kwmethod requires functions specify a keyword block.\nUse @kwmethod `f(args...;)` to specify no keywords.")
184219

185220
kwargs = fcall.args[2].args
186-
fargs = fcall.args[3:end]
221+
posargs = fcall.args[3:end]
222+
kwmethod_expr(f, posargs, kwargs, wherestack, body)
223+
end
224+
187225

226+
function kwmethod_expr(f, posargs, kwargs, wherestack, body)
188227
sort!(kwargs, by=argsym)
189228

190229
kwsyms = argsym.(kwargs)
@@ -196,11 +235,12 @@ macro kwmethod(fexpr)
196235
F = :(::($(esc(f)) isa Type ? Type{$(esc(f))} : typeof($(esc(f)))))
197236
end
198237

199-
fexpr.args[1] = :(KeywordDispatch.kwcall(($(esc.(kwsyms)...),)::NamedTuple{($(QuoteNode.(kwsyms)...),),T},
200-
$F,
201-
$(esc.(fargs)...)) where {T<:Tuple{$(esc.(kwtypes)...)}})
202-
return outexpr
238+
quote
239+
$(wrap_where(:(KeywordDispatch.kwcall(($(esc.(kwsyms)...),)::NamedTuple{($(QuoteNode.(kwsyms)...),),T},
240+
$F,
241+
$(esc.(posargs)...)) where {T<:Tuple{$(esc.(kwtypes)...)}}), wherestack)) =
242+
$body
243+
end
203244
end
204245

205-
206246
end # module

test/runtests.jl

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,3 +190,22 @@ end
190190
@test_throws MethodError Base.RoundDown()
191191
end
192192

193+
@testset "kwdispatch extra argument" begin
194+
@kwdispatch h(x) begin
195+
(a) -> x+a
196+
(b) -> x-b
197+
end
198+
199+
@test h(1, a=2) == 3
200+
@test h(1, b=2) == -1
201+
202+
@test_throws KeywordMethodError h(1)
203+
@test_throws KeywordMethodError h(1, c=2)
204+
205+
@test_throws MethodError h()
206+
@test_throws MethodError h(;a=1)
207+
208+
@kwmethod h(x; c) = c
209+
210+
@test h(1, c=2) == 2
211+
end

0 commit comments

Comments
 (0)