Merge pull request #2310 from SciML/myb/extend

Simplify `at extend`
SciML · Oct 9, 2023 · 563f04f · 563f04f
2 parents 298a9a5 + bdc45d1
commit 563f04f
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 23 deletions.
diff --git a/src/systems/model_parsing.jl b/src/systems/model_parsing.jl
@@ -372,6 +372,28 @@ function extend_args!(a, b, dict, expr, kwargs, varexpr, has_param = false)
     end
 end
 
+const EMPTY_DICT = Dict()
+const EMPTY_VoVoSYMBOL = Vector{Symbol}[]
+
+function Base.names(model::Model)
+    vars = keys(get(model.structure, :variables, EMPTY_DICT))
+    vars = union(vars, keys(get(model.structure, :parameters, EMPTY_DICT)))
+    vars = union(vars,
+        map(first, get(model.structure, :components, EMPTY_VoVoSYMBOL)))
+    collect(vars)
+end
+
+function _parse_extend!(ext, a, b, dict, expr, kwargs, varexpr, vars)
+    extend_args!(a, b, dict, expr, kwargs, varexpr)
+    ext[] = a
+    push!(b.args, Expr(:kw, :name, Meta.quot(a)))
+    push!(expr.args, :($a = $b))
+
+    dict[:extend] = [Symbol.(vars.args), a, b.args[1]]
+
+    push!(expr.args, :(@unpack $vars = $a))
+end
+
 function parse_extend!(exprs, ext, dict, mod, body, kwargs)
     expr = Expr(:block)
     varexpr = Expr(:block)
@@ -386,33 +408,27 @@ function parse_extend!(exprs, ext, dict, mod, body, kwargs)
                     error("`@extend` destructuring only takes an tuple as LHS. Got $body")
                 end
                 a, b = b.args
-            elseif Meta.isexpr(b, :call)
-                if (model = getproperty(mod, b.args[1])) isa Model
-                    _vars = keys(get(model.structure, :variables, Dict()))
-                    _vars = union(_vars, keys(get(model.structure, :parameters, Dict())))
-                    _vars = union(_vars,
-                        map(first, get(model.structure, :components, Vector{Symbol}[])))
-                    vars = Expr(:tuple)
-                    append!(vars.args, collect(_vars))
-                else
-                    error("Cannot infer the exact `Model` that `@extend $(body)` refers." *
-                          " Please specify the names that it brings into scope by:" *
-                          " `@extend a, b = oneport = OnePort()`.")
-                end
+                _parse_extend!(ext, a, b, dict, expr, kwargs, varexpr, vars)
+            else
+                error("When explicitly destructing in `@extend` please use the syntax: `@extend a, b = oneport = OnePort()`.")
             end
-            extend_args!(a, b, dict, expr, kwargs, varexpr)
-            ext[] = a
-            push!(b.args, Expr(:kw, :name, Meta.quot(a)))
-            push!(expr.args, :($a = $b))
-
-            dict[:extend] = [Symbol.(vars.args), a, b.args[1]]
-
-            if vars !== nothing
-                push!(expr.args, :(@unpack $vars = $a))
+        end
+        Expr(:call, a′, _...) => begin
+            a = Symbol(Symbol("#mtkmodel"), :__anonymous__, a′)
+            b = body
+            if (model = getproperty(mod, b.args[1])) isa Model
+                vars = Expr(:tuple)
+                append!(vars.args, names(model))
+                _parse_extend!(ext, a, b, dict, expr, kwargs, varexpr, vars)
+            else
+                error("Cannot infer the exact `Model` that `@extend $(body)` refers." *
+                      " Please specify the names that it brings into scope by:" *
+                      " `@extend a, b = oneport = OnePort()`.")
             end
         end
         _ => error("`@extend` only takes an assignment expression. Got $body")
     end
+    return nothing
 end
 
 function parse_variable_arg!(expr, vs, dict, mod, arg, varclass, kwargs)

diff --git a/test/model_parsing.jl b/test/model_parsing.jl
@@ -100,7 +100,7 @@ end
     @parameters begin
         C, [unit = u"F"]
     end
-    @extend oneport = OnePort(; v = 0.0)
+    @extend OnePort(; v = 0.0)
     @icon "https://upload.wikimedia.org/wikipedia/commons/7/78/Capacitor_symbol.svg"
     @equations begin
         D(v) ~ i / C