Fix models with type annotations (#183) (#184)

Fixes #183.

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.9.6"
+version = "0.9.7"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/compiler.jl

@@ -70,7 +70,7 @@ function model(expr, warn)
 
     # Generate main body
     modelinfo[:body] = generate_mainbody(
-        modelinfo[:modeldef][:body], modelinfo[:allargs_exprs], warn
+        modelinfo[:modeldef][:body], modelinfo[:allargs_syms], warn
     )
 
     return build_output(modelinfo)
@@ -117,7 +117,7 @@ function build_model_info(input_expr)
     allargs_exprs = first.(allargs_exprs_defaults)
 
     # Extract the names of the arguments.
-    allargs_syms = map(allargs_exprs_defaults) do (arg, _)
+    allargs_syms = map(allargs_exprs) do arg
         MacroTools.@match arg begin
             (::Type{T_}) | (name_::Type{T_}) => T
             name_::T_ => name

test/compiler.jl

@@ -635,4 +635,15 @@ end
         m = demo(0.)
         @test m() isa Float64
     end
+    @testset "type annotations" begin
+        @model function demo_without(x)
+            x ~ Normal()
+        end
+        @test isempty(VarInfo(demo_without(0.0)))
+
+        @model function demo_with(x::Real)
+            x ~ Normal()
+        end
+        @test isempty(VarInfo(demo_with(0.0)))
+    end
 end