fix: handle more typed array cases

This also fixes more tests
SciML · Sep 27, 2024 · 563dedc · 563dedc
1 parent ed5e0fb
commit 563dedc
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 17 deletions.
diff --git a/src/systems/model_parsing.jl b/src/systems/model_parsing.jl
@@ -232,55 +232,56 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
                 varclass, where_types, meta)
             return var, def, Dict()
         end
-        Expr(:tuple, Expr(:ref, a, b...), y) || Expr(:tuple, Expr(:(::), Expr(:ref, a, b...), n), y) => begin
-            isdefined(mod, :n) || (n = Real)
+        Expr(:tuple, Expr(:(::), Expr(:ref, a, b...), type), y) || Expr(:tuple, Expr(:ref, a, b...), y) => begin
+            (@isdefined type) || (type = Real)
             varname = Meta.isexpr(a, :call) ? a.args[1] : a
             push!(kwargs, Expr(:kw, varname, nothing))
             varval = unit_handled_variable_value(mod, y, varname)
             if varclass == :parameters
-                var = :($varname = $first(@parameters $a[$(b...)]::$n = ($varval, $y)))
+                var = :($varname = $first(@parameters $a[$(b...)]::$type = ($varval, $y)))
             else
-                var = :($varname = $first(@variables $a[$(b...)]::$n = ($varval, $y)))
+                var = :($varname = $first(@variables $a[$(b...)]::$type = ($varval, $y)))
             end
             #TODO: update `dict` aka `Model.structure` with the metadata
             (:($varname...), var), nothing, Dict()
         end
-        Expr(:(=), Expr(:ref, a, b...), y) || Expr(:(=), Expr(:(::), Expr(:ref, a, b...), n), y) => begin
-            isdefined(mod, :n) || (n = Real)
+        Expr(:(=), Expr(:(::), Expr(:ref, a, b...), type), y) || Expr(:(=), Expr(:ref, a, b...), y) => begin
+            (@isdefined type) || (type = Real)
             varname = Meta.isexpr(a, :call) ? a.args[1] : a
             if Meta.isexpr(y, :tuple)
                 varval = unit_handled_variable_value(mod, y, varname)
                 val, y = (y.args[1], y.args[2:end])
                 push!(kwargs, Expr(:kw, varname, nothing))
                 if varclass == :parameters
                     var = :($varname = $varname === nothing ? $val : $varname;
-                    $varname = $first(@parameters $a[$(b...)]::$n = (
+                    $varname = $first(@parameters $a[$(b...)]::$type = (
                         $varval, $(y...))))
                 else
                     var = :($varname = $varname === nothing ? $val : $varname;
-                    $varname = $first(@variables $a[$(b...)]::$n = (
+                    $varname = $first(@variables $a[$(b...)]::$type = (
                         $varval, $(y...))))
                 end
             else
                 push!(kwargs, Expr(:kw, varname, nothing))
                 if varclass == :parameters
                     var = :($varname = $varname === nothing ? $y : $varname;
-                    $varname = $first(@parameters $a[$(b...)]::$n = $varname))
+                    $varname = $first(@parameters $a[$(b...)]::$type = $varname))
                 else
                     var = :($varname = $varname === nothing ? $y : $varname;
-                    $varname = $first(@variables $a[$(b...)]::$n = $varname))
+                    $varname = $first(@variables $a[$(b...)]::$type = $varname))
                 end
             end
             #TODO: update `dict`` aka `Model.structure` with the metadata
             (:($varname...), var), nothing, Dict()
         end
-        Expr(:ref, a, b...) => begin
+        Expr(:(::), Expr(:ref, a, b...), type) || Expr(:ref, a, b...) => begin
+            (@isdefined type) || (type = Real)
             varname = a isa Expr && a.head == :call ? a.args[1] : a
             push!(kwargs, Expr(:kw, varname, nothing))
             if varclass == :parameters
-                var = :($varname = $first(@parameters $a[$(b...)] = $varname))
+                var = :($varname = $first(@parameters $a[$(b...)]::$type = $varname))
             elseif varclass == :variables
-                var = :($varname = $first(@variables $a[$(b...)] = $varname))
+                var = :($varname = $first(@variables $a[$(b...)]::$type = $varname))
             else
                 throw("Symbolic array with arbitrary length is not handled for $varclass.
                     Please open an issue with an example.")

diff --git a/test/model_parsing.jl b/test/model_parsing.jl
@@ -303,9 +303,9 @@ end
     @test symtype(type_model.par2) == Int
     @test symtype(type_model.par3) == BigFloat
     @test symtype(type_model.par4) == Float64
-    @test_broken symtype(type_model.par5[1]) == BigFloat
-    @test_broken symtype(type_model.par6[1]) == BigFloat
-    @test_broken symtype(type_model.par7[1, 1]) == BigFloat
+    @test symtype(type_model.par5[1]) == BigFloat
+    @test symtype(type_model.par6[1]) == BigFloat
+    @test symtype(type_model.par7[1, 1]) == BigFloat
 
     @test_throws TypeError TypeModel(; name = :throws, flag = 1)
     @test_throws TypeError TypeModel(; name = :throws, par0 = 1)
@@ -317,7 +317,7 @@ end
 
     # Test that array types are correctly added.
     @named type_model2 = TypeModel(; par5 = rand(BigFloat, 3))
-    @test_broken symtype(type_model2.par5[1]) == BigFloat
+    @test symtype(type_model2.par5[1]) == BigFloat
 
     @named type_model3 = TypeModel(; par7 = rand(BigFloat, 3, 3))
     @test symtype(type_model3.par7[1, 1]) == BigFloat