fix: pass type to var generation

Author: ven-k

docs/src/basics/MTKModel_Connector.md

@@ -246,7 +246,7 @@ Dict{Symbol, Any} with 7 entries:
   :components            => [[:model_a, :ModelA]]
   :variables             => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var), :v_array=>Dict(:size=>(2, 3)))
   :icon                  => URI("https://github.com/SciML/SciMLDocs/blob/main/docs/src/assets/logo.png")
-  :kwargs                => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin), :v=>Dict{Symbol, Union{Nothing, Symbol}}(:value=>:v_var, :type=>nothing), :v_array=>Dict(:value=>nothing, :type=>nothing), :p1=>Dict(:value=>nothing))
+  :kwargs                => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin), :v=>Dict{Symbol, Union{Nothing, Symbol}}(:value=>:v_var, :type=>Real), :v_array=>Dict(:value=>nothing, :type=>Real), :p1=>Dict(:value=>nothing))
   :structural_parameters => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin))
   :independent_variable  => t
   :extend                => Any[[:p2, :p1], Symbol("#mtkmodel__anonymous__ModelB"), :ModelB]

src/systems/model_parsing.jl

@@ -110,7 +110,7 @@ end
 
 function parse_variable_def!(dict, mod, arg, varclass, kwargs;
         def = nothing, indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
-        type::Union{Type, Nothing} = nothing)
+        type::Type = Real)
     metatypes = [(:connection_type, VariableConnectType),
         (:description, VariableDescription),
         (:unit, VariableUnit),
@@ -133,7 +133,7 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs;
             else
                 push!(kwargs, Expr(:kw, Expr(:(::), a, Union{Nothing, type}), nothing))
             end
-            var = generate_var!(dict, a, varclass; indices)
+            var = generate_var!(dict, a, varclass; indices, type)
             dict[:kwargs][getname(var)] = Dict(:value => def, :type => type)
             (var, def)
         end
@@ -153,15 +153,15 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs;
             else
                 push!(kwargs, Expr(:kw, Expr(:(::), a, Union{Nothing, type}), nothing))
             end
-            var = generate_var!(dict, a, b, varclass; indices)
+            var = generate_var!(dict, a, b, varclass; indices, type)
             type !== nothing && (dict[varclass][getname(var)][:type] = type)
             dict[:kwargs][getname(var)] = Dict(:value => def, :type => type)
             (var, def)
         end
         Expr(:(=), a, b) => begin
             Base.remove_linenums!(b)
             def, meta = parse_default(mod, b)
-            var, def = parse_variable_def!(dict, mod, a, varclass, kwargs; def)
+            var, def = parse_variable_def!(dict, mod, a, varclass, kwargs; def, type)
             dict[varclass][getname(var)][:default] = def
             if meta !== nothing
                 for (type, key) in metatypes
@@ -179,7 +179,7 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs;
             (var, def)
         end
         Expr(:tuple, a, b) => begin
-            var, def = parse_variable_def!(dict, mod, a, varclass, kwargs)
+            var, def = parse_variable_def!(dict, mod, a, varclass, kwargs; type)
             meta = parse_metadata(mod, b)
             if meta !== nothing
                 for (type, key) in metatypes
@@ -200,34 +200,39 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs;
         Expr(:ref, a, b...) => begin
             indices = map(i -> UnitRange(i.args[2], i.args[end]), b)
             parse_variable_def!(dict, mod, a, varclass, kwargs;
-                def, indices)
+                def, indices, type)
         end
         _ => error("$arg cannot be parsed")
     end
 end
 
 function generate_var(a, varclass;
-        indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing)
-    var = indices === nothing ? Symbolics.variable(a) : first(@variables $a[indices...])
+        indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
+        type = Real)
+    var = indices === nothing ? Symbolics.variable(a; T = type) :
+          first(@variables $a[indices...]::type)
     if varclass == :parameters
         var = toparam(var)
     end
     var
 end
 
 function generate_var!(dict, a, varclass;
-        indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing)
+        indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
+        type = Real)
     vd = get!(dict, varclass) do
         Dict{Symbol, Dict{Symbol, Any}}()
     end
     vd isa Vector && (vd = first(vd))
     vd[a] = Dict{Symbol, Any}()
     indices !== nothing && (vd[a][:size] = Tuple(lastindex.(indices)))
-    generate_var(a, varclass; indices)
+    generate_var(a, varclass; indices, type)
 end
 
 function generate_var!(dict, a, b, varclass;
-        indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing)
+        indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
+        type = Real)
+    # (type isa Nothing && type = Real)
     iv = generate_var(b, :variables)
     prev_iv = get!(dict, :independent_variable) do
         iv
@@ -239,10 +244,10 @@ function generate_var!(dict, a, b, varclass;
     vd isa Vector && (vd = first(vd))
     vd[a] = Dict{Symbol, Any}()
     var = if indices === nothing
-        Symbolics.variable(a, T = SymbolicUtils.FnType{Tuple{Any}, Real})(iv)
+        Symbolics.variable(a, T = SymbolicUtils.FnType{Tuple{Any}, type})(iv)
     else
         vd[a][:size] = Tuple(lastindex.(indices))
-        first(@variables $a(iv)[indices...])
+        first(@variables $a(iv)[indices...]::type)
     end
     if varclass == :parameters
         var = toparam(var)

test/model_parsing.jl

@@ -1,7 +1,7 @@
 using ModelingToolkit, Test
 using ModelingToolkit: get_gui_metadata, get_systems, get_connector_type,
-                       get_ps, getdefault, getname, scalarize, VariableDescription,
-                       RegularConnector
+                       get_ps, getdefault, getname, scalarize, symtype,
+                       VariableDescription, RegularConnector
 using URIs: URI
 using Distributions
 using DynamicQuantities, OrdinaryDiffEq
@@ -251,14 +251,23 @@ end
             par1::Int = 1
             par2(t)::Int,
             [description = "Enforced `par4` to be an Int by setting the type to the keyword-arg."]
-            par3(t)::Float64 = 1.0
+            par3(t)::BigFloat = 1.0
             par4(t)::Float64 = 1 # converts 1 to 1.0 of Float64 type
+            par5[1:3]::BigFloat
+            par6(t)[1:3]::BigFloat
+            par7(t)[1:3]::BigFloat = 1.0, [description = "with description"]
         end
     end
 
     @named type_model = TypeModel()
 
-    @test getname.(parameters(type_model)) == [:par0, :par1, :par2, :par3, :par4]
+    @test symtype(type_model.par1) == Int
+    @test symtype(type_model.par2) == Int
+    @test symtype(type_model.par3) == BigFloat
+    @test symtype(type_model.par4) == Float64
+    @test symtype(type_model.par5[1]) == BigFloat
+    @test symtype(type_model.par6[1]) == BigFloat
+    @test symtype(type_model.par7[1]) == BigFloat
 
     @test_throws TypeError TypeModel(; name = :throws, flag = 1)
     @test_throws TypeError TypeModel(; name = :throws, par0 = 1)
@@ -350,8 +359,8 @@ end
     @test A.structure[:extend] == [[:e], :extended_e, :E]
     @test A.structure[:equations] == ["e ~ 0"]
     @test A.structure[:kwargs] ==
-          Dict{Symbol, Dict}(:p => Dict(:value => nothing, :type => nothing),
-        :v => Dict(:value => nothing, :type => nothing))
+          Dict{Symbol, Dict}(:p => Dict(:value => nothing, :type => Real),
+        :v => Dict(:value => nothing, :type => Real))
     @test A.structure[:components] == [[:cc, :C]]
 end