Merge pull request #2306 from SciML/myb/extend

Support implicit name unpack in `at extend`

Author: YingboMa

src/systems/model_parsing.jl

@@ -237,7 +237,7 @@ function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps,
     if mname == Symbol("@components")
         parse_components!(exprs, comps, dict, body, kwargs)
     elseif mname == Symbol("@extend")
-        parse_extend!(exprs, ext, dict, body, kwargs)
+        parse_extend!(exprs, ext, dict, mod, body, kwargs)
     elseif mname == Symbol("@variables")
         parse_variables!(exprs, vs, dict, mod, body, :variables, kwargs)
     elseif mname == Symbol("@parameters")
@@ -372,24 +372,35 @@ function extend_args!(a, b, dict, expr, kwargs, varexpr, has_param = false)
     end
 end
 
-function parse_extend!(exprs, ext, dict, body, kwargs)
+function parse_extend!(exprs, ext, dict, mod, body, kwargs)
     expr = Expr(:block)
     varexpr = Expr(:block)
     push!(exprs, varexpr)
     push!(exprs, expr)
     body = deepcopy(body)
     MLStyle.@match body begin
         Expr(:(=), a, b) => begin
-            vars = nothing
             if Meta.isexpr(b, :(=))
                 vars = a
                 if !Meta.isexpr(vars, :tuple)
                     error("`@extend` destructuring only takes an tuple as LHS. Got $body")
                 end
                 a, b = b.args
-                extend_args!(a, b, dict, expr, kwargs, varexpr)
-                vars, a, b
+            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
             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))

test/model_parsing.jl

@@ -7,77 +7,76 @@ using Unitful
 
 ENV["MTK_ICONS_DIR"] = "$(@__DIR__)/icons"
 
-@testset "Comprehensive Test of Parsing Models (with an RC Circuit)" begin
-    @connector RealInput begin
-        u(t), [input = true, unit = u"V"]
+@connector RealInput begin
+    u(t), [input = true, unit = u"V"]
+end
+@connector RealOutput begin
+    u(t), [output = true, unit = u"V"]
+end
+@mtkmodel Constant begin
+    @components begin
+        output = RealOutput()
     end
-    @connector RealOutput begin
-        u(t), [output = true, unit = u"V"]
+    @parameters begin
+        k, [description = "Constant output value of block"]
     end
-    @mtkmodel Constant begin
-        @components begin
-            output = RealOutput()
-        end
-        @parameters begin
-            k, [description = "Constant output value of block"]
-        end
-        @equations begin
-            output.u ~ k
-        end
+    @equations begin
+        output.u ~ k
     end
+end
 
-    @variables t [unit = u"s"]
-    D = Differential(t)
+@variables t [unit = u"s"]
+D = Differential(t)
 
-    @connector Pin begin
-        v(t), [unit = u"V"]                    # Potential at the pin [V]
-        i(t), [connect = Flow, unit = u"A"]    # Current flowing into the pin [A]
-        @icon "pin.png"
-    end
+@connector Pin begin
+    v(t), [unit = u"V"]                    # Potential at the pin [V]
+    i(t), [connect = Flow, unit = u"A"]    # Current flowing into the pin [A]
+    @icon "pin.png"
+end
 
-    @named p = Pin(; v = π)
-    @test getdefault(p.v) == π
-    @test Pin.isconnector == true
+@named p = Pin(; v = π)
+@test getdefault(p.v) == π
+@test Pin.isconnector == true
 
-    @mtkmodel OnePort begin
-        @components begin
-            p = Pin()
-            n = Pin()
-        end
-        @variables begin
-            v(t), [unit = u"V"]
-            i(t), [unit = u"A"]
-        end
-        @icon "oneport.png"
-        @equations begin
-            v ~ p.v - n.v
-            0 ~ p.i + n.i
-            i ~ p.i
-        end
+@mtkmodel OnePort begin
+    @components begin
+        p = Pin()
+        n = Pin()
+    end
+    @variables begin
+        v(t), [unit = u"V"]
+        i(t), [unit = u"A"]
     end
+    @icon "oneport.png"
+    @equations begin
+        v ~ p.v - n.v
+        0 ~ p.i + n.i
+        i ~ p.i
+    end
+end
 
-    @test OnePort.isconnector == false
+@test OnePort.isconnector == false
 
-    @mtkmodel Ground begin
-        @components begin
-            g = Pin()
-        end
-        @icon begin
-            read(abspath(ENV["MTK_ICONS_DIR"], "ground.svg"), String)
-        end
-        @equations begin
-            g.v ~ 0
-        end
+@mtkmodel Ground begin
+    @components begin
+        g = Pin()
+    end
+    @icon begin
+        read(abspath(ENV["MTK_ICONS_DIR"], "ground.svg"), String)
     end
+    @equations begin
+        g.v ~ 0
+    end
+end
 
-    resistor_log = "$(@__DIR__)/logo/resistor.svg"
-    @mtkmodel Resistor begin
-        @extend v, i = oneport = OnePort()
-        @parameters begin
-            R, [unit = u"Ω"]
-        end
-        @icon begin
-            """<?xml version="1.0" encoding="UTF-8"?>
+resistor_log = "$(@__DIR__)/logo/resistor.svg"
+@mtkmodel Resistor begin
+    @extend v, i = oneport = OnePort()
+    @parameters begin
+        R, [unit = u"Ω"]
+    end
+    @icon begin
+        """<?xml version="1.0" encoding="UTF-8"?>
 <svg xmlns="http://www.w3.org/2000/svg" width="80" height="30">
   <path d="M10 15
 l15 0
@@ -91,88 +90,87 @@ l2.5 -5
 l15 0" stroke="black" stroke-width="1" stroke-linejoin="bevel" fill="none"></path>
 </svg>
 """
-        end
-        @equations begin
-            v ~ i * R
-        end
     end
+    @equations begin
+        v ~ i * R
+    end
+end
 
-    @mtkmodel Capacitor begin
-        @parameters begin
-            C, [unit = u"F"]
-        end
-        @extend v, i = oneport = OnePort(; v = 0.0)
-        @icon "https://upload.wikimedia.org/wikipedia/commons/7/78/Capacitor_symbol.svg"
-        @equations begin
-            D(v) ~ i / C
-        end
+@mtkmodel Capacitor begin
+    @parameters begin
+        C, [unit = u"F"]
+    end
+    @extend oneport = OnePort(; v = 0.0)
+    @icon "https://upload.wikimedia.org/wikipedia/commons/7/78/Capacitor_symbol.svg"
+    @equations begin
+        D(v) ~ i / C
     end
+end
 
-    @named capacitor = Capacitor(C = 10, v = 10.0)
-    @test getdefault(capacitor.v) == 10.0
+@named capacitor = Capacitor(C = 10, v = 10.0)
+@test getdefault(capacitor.v) == 10.0
 
-    @mtkmodel Voltage begin
-        @extend v, i = oneport = OnePort()
-        @components begin
-            V = RealInput()
-        end
-        @equations begin
-            v ~ V.u
-        end
+@mtkmodel Voltage begin
+    @extend v, i = oneport = OnePort()
+    @components begin
+        V = RealInput()
     end
+    @equations begin
+        v ~ V.u
+    end
+end
 
-    @mtkmodel RC begin
-        @structural_parameters begin
-            R_val = 10
-            C_val = 10
-            k_val = 10
-        end
-        @components begin
-            resistor = Resistor(; R = R_val)
-            capacitor = Capacitor(; C = C_val)
-            source = Voltage()
-            constant = Constant(; k = k_val)
-            ground = Ground()
-        end
-
-        @equations begin
-            connect(constant.output, source.V)
-            connect(source.p, resistor.p)
-            connect(resistor.n, capacitor.p)
-            connect(capacitor.n, source.n, ground.g)
-        end
+@mtkmodel RC begin
+    @structural_parameters begin
+        R_val = 10
+        C_val = 10
+        k_val = 10
+    end
+    @components begin
+        resistor = Resistor(; R = R_val)
+        capacitor = Capacitor(; C = C_val)
+        source = Voltage()
+        constant = Constant(; k = k_val)
+        ground = Ground()
     end
 
-    C_val = 20
-    R_val = 20
-    res__R = 100
-    @mtkbuild rc = RC(; C_val, R_val, resistor.R = res__R)
-    resistor = getproperty(rc, :resistor; namespace = false)
-    @test getname(rc.resistor) === getname(resistor)
-    @test getname(rc.resistor.R) === getname(resistor.R)
-    @test getname(rc.resistor.v) === getname(resistor.v)
-    # Test that `resistor.R` overrides `R_val` in the argument.
-    @test getdefault(rc.resistor.R) == res__R != R_val
-    # Test that `C_val` passed via argument is set as default of C.
-    @test getdefault(rc.capacitor.C) == C_val
-    # Test that `k`'s default value is unchanged.
-    @test getdefault(rc.constant.k) == RC.structure[:kwargs][:k_val]
-    @test getdefault(rc.capacitor.v) == 0.0
-
-    @test get_gui_metadata(rc.resistor).layout == Resistor.structure[:icon] ==
-          read(joinpath(ENV["MTK_ICONS_DIR"], "resistor.svg"), String)
-    @test get_gui_metadata(rc.ground).layout ==
-          read(abspath(ENV["MTK_ICONS_DIR"], "ground.svg"), String)
-    @test get_gui_metadata(rc.capacitor).layout ==
-          URI("https://upload.wikimedia.org/wikipedia/commons/7/78/Capacitor_symbol.svg")
-    @test OnePort.structure[:icon] ==
-          URI("file:///" * abspath(ENV["MTK_ICONS_DIR"], "oneport.png"))
-    @test ModelingToolkit.get_gui_metadata(rc.resistor.p).layout == Pin.structure[:icon] ==
-          URI("file:///" * abspath(ENV["MTK_ICONS_DIR"], "pin.png"))
-
-    @test length(equations(rc)) == 1
+    @equations begin
+        connect(constant.output, source.V)
+        connect(source.p, resistor.p)
+        connect(resistor.n, capacitor.p)
+        connect(capacitor.n, source.n, ground.g)
+    end
 end
 
+C_val = 20
+R_val = 20
+res__R = 100
+@mtkbuild rc = RC(; C_val, R_val, resistor.R = res__R)
+resistor = getproperty(rc, :resistor; namespace = false)
+@test getname(rc.resistor) === getname(resistor)
+@test getname(rc.resistor.R) === getname(resistor.R)
+@test getname(rc.resistor.v) === getname(resistor.v)
+# Test that `resistor.R` overrides `R_val` in the argument.
+@test getdefault(rc.resistor.R) == res__R != R_val
+# Test that `C_val` passed via argument is set as default of C.
+@test getdefault(rc.capacitor.C) == C_val
+# Test that `k`'s default value is unchanged.
+@test getdefault(rc.constant.k) == RC.structure[:kwargs][:k_val]
+@test getdefault(rc.capacitor.v) == 0.0
+
+@test get_gui_metadata(rc.resistor).layout == Resistor.structure[:icon] ==
+      read(joinpath(ENV["MTK_ICONS_DIR"], "resistor.svg"), String)
+@test get_gui_metadata(rc.ground).layout ==
+      read(abspath(ENV["MTK_ICONS_DIR"], "ground.svg"), String)
+@test get_gui_metadata(rc.capacitor).layout ==
+      URI("https://upload.wikimedia.org/wikipedia/commons/7/78/Capacitor_symbol.svg")
+@test OnePort.structure[:icon] ==
+      URI("file:///" * abspath(ENV["MTK_ICONS_DIR"], "oneport.png"))
+@test ModelingToolkit.get_gui_metadata(rc.resistor.p).layout == Pin.structure[:icon] ==
+      URI("file:///" * abspath(ENV["MTK_ICONS_DIR"], "pin.png"))
+
+@test length(equations(rc)) == 1
+
 @testset "Parameters and Structural parameters in various modes" begin
     @mtkmodel MockModel begin
         @parameters begin