Skip to content

Commit 879383a

Browse files
committed
feat: explicitly specify the model kwargs and args
1 parent 4e3cb7a commit 879383a

File tree

1 file changed

+25
-9
lines changed

1 file changed

+25
-9
lines changed

src/systems/model_parsing.jl

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -79,17 +79,17 @@ function parse_variables_with_kw!(exprs, var, dict, mod, body, varexpr, varclass
7979
Expr(:(=), a, b::Symbol) => begin
8080
isdefined(mod, b) ?
8181
parse_variables!(exprs, var, dict, mod, arg, varclass, kwargs) :
82-
push!(kwargs, b)
82+
push!(varexpr.args[end].args[end].args, arg)
8383
end
8484
Expr(:(=), a, b) => begin
85-
def = Base.remove_linenums!(b).args[end]
85+
def = Base.remove_linenums!(b)
8686
MLStyle.@match def begin
8787
Expr(:tuple, x::Symbol, y) || x::Symbol => begin
88-
push!(varexpr.args[end].args[end].args, :($a = $def))
89-
push!(kwargs, x)
88+
push!(varexpr.args[end].args[end].args, :($a = $(def.args[end])))
9089
end
91-
Expr(:tuple, x::Number, y) => parse_variables!(exprs, var, dict, mod, arg, varclass, kwargs)
90+
Expr(:tuple, x::Number, y) => (@info "111"; parse_variables!(exprs, var, dict, mod, arg, varclass, kwargs))
9291
::Number => parse_variables!(exprs, var, dict, mod, arg, varclass, kwargs)
92+
::Expr => push!(varexpr.args[end].args[end].args, :($a = $(def.args[end])))
9393
_ => @info "Got $def"
9494
end
9595
end
@@ -166,7 +166,21 @@ macro model(name::Symbol, expr)
166166
esc(model_macro(__module__, name, expr))
167167
end
168168

169-
function model_macro(mod, name, expr)
169+
@inline is_kwarg(::Symbol) = false
170+
@inline is_kwarg(e::Expr) = (e.head == :parameters)
171+
172+
macro model(fcall::Expr, expr)
173+
fcall.head == :call || "Couldn't comprehend the model $arg"
174+
175+
arglist, kwargs = if lastindex(fcall.args) > 1 && is_kwarg(fcall.args[2])
176+
(lastindex(fcall.args) > 2 ? (@info 1; Set(fcall.args[3:end])) : (@info 2; Set())), Set(fcall.args[2].args)
177+
else
178+
Set(), Set(fcall.args[2:end])
179+
end
180+
esc(model_macro(__module__, fcall.args[1], expr; arglist, kwargs))
181+
end
182+
183+
function model_macro(mod, name, expr; arglist = Set([]), kwargs = Set([]))
170184
exprs = Expr(:block)
171185
dict = Dict{Symbol, Any}()
172186
comps = Symbol[]
@@ -179,7 +193,7 @@ function model_macro(mod, name, expr)
179193
end)
180194
varexpr = :(vss = @variables begin
181195
end)
182-
kwargs = []
196+
183197
for arg in expr.args
184198
arg isa LineNumberNode && continue
185199
if arg.head == :macrocall
@@ -208,7 +222,8 @@ function model_macro(mod, name, expr)
208222
else
209223
push!(exprs.args, :($extend($sys, $(ext[]))))
210224
end
211-
:($name = $Model((; name, $(kwargs...)) -> $exprs, $dict))
225+
226+
:($name = $Model(($(arglist...); name, $(kwargs...)) -> $exprs, $dict))
212227
end
213228

214229
function parse_model!(exprs, comps, ext, eqs, icon, vs, varexpr, ps, parexpr, dict, mod, arg, kwargs)
@@ -261,7 +276,8 @@ function _rename(compname, varname)
261276
end
262277

263278
function component_args!(a, b, expr, kwargs)
264-
for i in 2:lastindex(b.args)
279+
start = b.head == :parameters ? 1 : 2
280+
for i in start:lastindex(b.args)
265281
arg = b.args[i]
266282
arg isa LineNumberNode && continue
267283
MLStyle.@match arg begin

0 commit comments

Comments
 (0)