@@ -51,6 +51,22 @@ argtype(x::Symbol) = Any
51
51
argtype (x:: Expr ) = x. head == :(:: ) ? x. args[end ] : error (" unexpected expression $x " )
52
52
53
53
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
+
54
70
struct KeywordMethodError <: Exception
55
71
f
56
72
args
88
104
89
105
90
106
"""
91
- @kwdispatch expr
107
+ @kwdispatch sig [methods]
92
108
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
94
110
function can also be provided, in which case _all_ calls to that function are will
95
111
dispatch to keyword methods.
96
112
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.
99
117
100
118
# Examples
101
119
@@ -105,17 +123,18 @@ methods, use the [`@kwmethod`](@ref) macro.
105
123
106
124
@kwdispacth f # equivalent to @kwdispatch f(_...)
107
125
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
108
134
```
109
135
"""
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)
119
138
120
139
# handle: `fun`, `Mod.fun`, `(a::B)`
121
140
if fcall isa Symbol || fcall isa Expr && (fcall. head in (:., :(:: ), :curly ))
@@ -125,8 +144,8 @@ macro kwdispatch(fexpr)
125
144
@assert fcall isa Expr && fcall. head == :call
126
145
127
146
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
130
149
error (" keyword arguments should only appear in @kwdispatch expressions" )
131
150
end
132
151
f = argmeth (f)
@@ -137,13 +156,35 @@ macro kwdispatch(fexpr)
137
156
ftype = :(typeof ($ (esc (f))))
138
157
end
139
158
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 ))
142
161
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
145
167
end
146
168
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
+
147
188
"""
148
189
@kwmethod expr
149
190
@@ -165,16 +206,10 @@ The positional signature should first be designated by the [`@kwdispatch`](@ref)
165
206
"""
166
207
macro kwmethod (fexpr)
167
208
@assert fexpr isa Expr && fexpr. head in (:function , :(= ))
168
- fexpr . args[ 2 ] = esc (fexpr. args[2 ])
209
+ body = esc (fexpr. args[2 ])
169
210
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 ])
176
212
177
- fcall = fexpr. args[1 ]
178
213
@assert fcall isa Expr && fcall. head == :call
179
214
180
215
f = fcall. args[1 ]
@@ -183,8 +218,12 @@ macro kwmethod(fexpr)
183
218
error (" @kwmethod requires functions specify a keyword block.\n Use @kwmethod `f(args...;)` to specify no keywords." )
184
219
185
220
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
+
187
225
226
+ function kwmethod_expr (f, posargs, kwargs, wherestack, body)
188
227
sort! (kwargs, by= argsym)
189
228
190
229
kwsyms = argsym .(kwargs)
@@ -196,11 +235,12 @@ macro kwmethod(fexpr)
196
235
F = :(:: ($(esc(f) ) isa Type ? Type{$ (esc (f))} : typeof ($ (esc (f)))))
197
236
end
198
237
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
203
244
end
204
245
205
-
206
246
end # module
0 commit comments