Merge pull request #99 from TuringLang/tor/doc-string-for-model

TuringLang · May 3, 2020 · 0758d01 · 0758d01
2 parents 0ec01ae + acb21a4
commit 0758d01
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/src/compiler.jl b/src/compiler.jl
@@ -352,7 +352,7 @@ function build_output(model_info)
         $generator($(args...)) = $(DynamicPPL.Model)($evaluator, $args_nt, $model_gen_constructor)
         $(generator_kw_form...)
 
-        $model_gen = $model_gen_constructor
+        $(Base).@__doc__ $model_gen = $model_gen_constructor
     end
 end
 

diff --git a/test/compiler.jl b/test/compiler.jl
@@ -615,4 +615,19 @@ end
 
         @test lp_w_threads ≈ lp_wo_threads
     end
+
+    @testset "docstring" begin
+        "This is a test"
+        @model function demo(x)
+            m ~ Normal()
+            x ~ Normal(m, 1)
+        end
+
+        s = @doc(demo)
+        @test string(s) == "This is a test\n"
+
+        # Verify that adding docstring didn't completely break execution of model
+        m = demo(0.)
+        @test m() isa Float64
+    end
 end