feat: include all of variable metadata in the Model.structure

SciML · Jul 5, 2023 · 013f363 · 013f363
1 parent a1a02a2
commit 013f363
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 4 deletions.
diff --git a/src/systems/model_parsing.jl b/src/systems/model_parsing.jl
@@ -60,6 +60,22 @@ function connector_macro(mod, name, body)
 end
 
 function parse_variable_def!(dict, mod, arg, varclass, kwargs, def = nothing)
+    metatypes = [(:connection_type, VariableConnectType),
+        (:description, VariableDescription),
+        (:unit, VariableUnit),
+        (:bounds, VariableBounds),
+        (:noise, VariableNoiseType),
+        (:input, VariableInput),
+        (:output, VariableOutput),
+        (:irreducible, VariableIrreducible),
+        (:state_priority, VariableStatePriority),
+        (:misc, VariableMisc),
+        (:disturbance, VariableDisturbance),
+        (:tunable, VariableTunable),
+        (:dist, VariableDistribution),
+        (:binary, VariableBinary),
+        (:integer, VariableInteger)]
+
     arg isa LineNumberNode && return
     MLStyle.@match arg begin
         a::Symbol => begin
@@ -80,8 +96,11 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, def = nothing)
             var, _ = parse_variable_def!(dict, mod, a, varclass, kwargs, def)
             dict[varclass][getname(var)][:default] = def
             if !isnothing(meta)
-                if (ct = get(meta, VariableConnectType, nothing)) !== nothing
-                    dict[varclass][getname(var)][:connection_type] = nameof(ct)
+                for (type, key) in metatypes
+                    if (mt = get(meta, key, nothing)) !== nothing
+                        key == VariableConnectType && (mt = nameof(mt))
+                        dict[varclass][getname(var)][type] = mt
+                    end
                 end
                 var = set_var_metadata(var, meta)
             end
@@ -90,8 +109,14 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, def = nothing)
         Expr(:tuple, a, b) => begin
             var, def = parse_variable_def!(dict, mod, a, varclass, kwargs)
             meta = parse_metadata(mod, b)
-            if (ct = get(meta, VariableConnectType, nothing)) !== nothing
-                dict[varclass][getname(var)][:connection_type] = nameof(ct)
+            if !isnothing(meta)
+                for (type, key) in metatypes
+                    if (mt = get(meta, key, nothing)) !== nothing
+                        key == VariableConnectType && (mt = nameof(mt))
+                        dict[varclass][getname(var)][type] = mt
+                    end
+                end
+                var = set_var_metadata(var, meta)
             end
             (set_var_metadata(var, meta), def)
         end

diff --git a/test/model_parsing.jl b/test/model_parsing.jl
@@ -216,3 +216,20 @@ getdefault(a.b.k) == 1
 getdefault(a.b.i) == 20
 getdefault(a.b.j) == 30
 getdefault(a.b.k) == 40
+
+metadata = Dict(:description => "Variable to test metadata in the Model.structure",
+    :input => true, :bounds => :((-1, 1)), :connection_type => :Flow, :integer => true,
+    :binary => false, :tunable => false, :disturbance => true, :dist => :(Normal(1, 1)))
+
+@connector MockMeta begin
+    function m(t)
+        1,
+        [description = "Variable to test metadata in the Model.structure",
+            input = true, bounds = (-1, 1), connect = Flow, integer = true,
+            binary = false, tunable = false, disturbance = true, dist = Normal(1, 1)]
+    end
+end
+
+for (k, v) in metadata
+    @test MockMeta.structure[:variables][:m][k] == v
+end