Merge pull request #168 from ReactiveBayes/165-function-based-syntax-…

…for-meta-and-constraints-macros-parses-but-does-produce-weird-errors-when-applied 165 function based syntax for meta and constraints macros parses but does produce weird errors when applied
ReactiveBayes · Mar 12, 2024 · 42a1ba8 · 42a1ba8
2 parents 39d7b2a + c06d375
commit 42a1ba8
Show file tree

Hide file tree

Showing 4 changed files with 163 additions and 8 deletions.
diff --git a/src/plugins/meta/meta_macro.jl b/src/plugins/meta/meta_macro.jl
@@ -17,10 +17,25 @@ and assigns it to the `__meta__` variable. It then evaluates the given expressio
 - `e::Expr`: The expression that will generate the `GraphPPL.MetaSpecification` object.
 """
 function add_meta_construction(e::Expr)
-    return quote
-        __meta__ = GraphPPL.MetaSpecification()
-        $e
-        __meta__
+    if @capture(e, (function m_name_(m_args__; m_kwargs__)
+        c_body_
+    end) | (function m_name_(m_args__)
+    c_body_
+end))
+    m_kwargs = m_kwargs === nothing ? [] : m_kwargs
+        return quote
+            function $m_name($(m_args...); $(m_kwargs...))
+                __meta__ = GraphPPL.MetaSpecification()
+                $c_body
+                return __meta__
+            end
+        end
+    else
+        return quote
+            __meta__ = GraphPPL.MetaSpecification()
+            $e
+            __meta__
+        end
     end
 end
 

diff --git a/src/plugins/variational_constraints/variational_constraints_macro.jl b/src/plugins/variational_constraints/variational_constraints_macro.jl
@@ -18,10 +18,23 @@ end
 check_for_returns_constraints = (x) -> check_for_returns(x; tag = "constraints")
 
 function add_constraints_construction(e::Expr)
-    return quote
-        __constraints__ = GraphPPL.Constraints()
-        $e
-        __constraints__
+    if @capture(e, (function c_name_(c_args__; c_kwargs__)
+        c_body_
+    end) | (function c_name_(c_args__) c_body_ end)) 
+        c_kwargs = c_kwargs === nothing ? [] : c_kwargs
+        return quote
+            function $c_name($(c_args...); $(c_kwargs...))
+                __constraints__ = GraphPPL.Constraints()
+                $c_body
+                return __constraints__
+            end
+        end
+    else
+        return quote
+            __constraints__ = GraphPPL.Constraints()
+            $e
+            __constraints__
+        end
     end
 end
 

diff --git a/test/plugins/meta/meta_macro_tests.jl b/test/plugins/meta/meta_macro_tests.jl
@@ -49,6 +49,82 @@ end
         __meta__
     end
     @test_expression_generating add_meta_construction(input) output
+
+    # Test 3: add_constraints_construction to constraint specification with function specification
+    input = quote 
+        function somemeta()
+            GCV(x, k, w) -> GCVMetadata(GaussHermiteCubature(20))
+            NormalMeanVariance() -> MyCustomMetaObject(arg1, arg2)
+            x -> MySecondCustomMetaObject(arg3)
+        end
+    end
+    output = quote 
+        function somemeta(;)
+            __meta__ = GraphPPL.MetaSpecification()
+            GCV(x, k, w) -> GCVMetadata(GaussHermiteCubature(20))
+            NormalMeanVariance() -> MyCustomMetaObject(arg1, arg2)
+            x -> MySecondCustomMetaObject(arg3)
+            return __meta__
+        end
+    end
+    @test_expression_generating add_meta_construction(input) output
+
+    # Test 4: add_constraints_construction to constraint specification with function specification and arguments
+    input = quote 
+        function somemeta(x, y)
+            GCV(x, k, w) -> GCVMetadata(GaussHermiteCubature(20))
+            NormalMeanVariance() -> MyCustomMetaObject(arg1, arg2)
+            x -> MySecondCustomMetaObject(arg3)
+        end
+    end
+    output = quote 
+        function somemeta(x, y;)
+            __meta__ = GraphPPL.MetaSpecification()
+            GCV(x, k, w) -> GCVMetadata(GaussHermiteCubature(20))
+            NormalMeanVariance() -> MyCustomMetaObject(arg1, arg2)
+            x -> MySecondCustomMetaObject(arg3)
+            return __meta__
+        end
+    end
+    @test_expression_generating add_meta_construction(input) output
+
+    # Test 5: add_constraints_construction to constraint specification with function specification and arguments and keyword arguments
+    input = quote 
+        function somemeta(x, y; z)
+            GCV(x, k, w) -> GCVMetadata(GaussHermiteCubature(20))
+            NormalMeanVariance() -> MyCustomMetaObject(arg1, arg2)
+            x -> MySecondCustomMetaObject(arg3)
+        end
+    end
+    output = quote 
+        function somemeta(x, y; z)
+            __meta__ = GraphPPL.MetaSpecification()
+            GCV(x, k, w) -> GCVMetadata(GaussHermiteCubature(20))
+            NormalMeanVariance() -> MyCustomMetaObject(arg1, arg2)
+            x -> MySecondCustomMetaObject(arg3)
+            return __meta__
+        end
+    end
+    @test_expression_generating add_meta_construction(input) output
+
+    # Test 6: add_constraints_construction to constraint specification with function specification and only keyword arguments
+    input = quote 
+        function somemeta(; z)
+            GCV(x, k, w) -> GCVMetadata(GaussHermiteCubature(20))
+            NormalMeanVariance() -> MyCustomMetaObject(arg1, arg2)
+            x -> MySecondCustomMetaObject(arg3)
+        end
+    end
+    output = quote 
+        function somemeta(; z)
+            __meta__ = GraphPPL.MetaSpecification()
+            GCV(x, k, w) -> GCVMetadata(GaussHermiteCubature(20))
+            NormalMeanVariance() -> MyCustomMetaObject(arg1, arg2)
+            x -> MySecondCustomMetaObject(arg3)
+            return __meta__
+        end
+    end
+    @test_expression_generating add_meta_construction(input) output
 end
 
 @testitem "create_submodel_meta" begin

diff --git a/test/plugins/variational_constraints/variational_constraints_macro_tests.jl b/test/plugins/variational_constraints/variational_constraints_macro_tests.jl
@@ -98,6 +98,57 @@ end
         __constraints__
     end
     @test_expression_generating add_constraints_construction(input) output
+
+    # Test 3: add_constraints_construction to constraint specification with function specification
+    input = quote 
+        function some_constraints()
+            q(x, y) = q(x)q(y)
+            q(x)::PointMass
+        end
+    end
+    output = quote 
+        function some_constraints(;)
+            __constraints__ = GraphPPL.Constraints()
+            q(x, y) = q(x)q(y)
+            q(x)::PointMass
+            return __constraints__
+        end
+    end
+    @test_expression_generating add_constraints_construction(input) output
+
+    # Test 4: add_constraints_construction to constraint specification with function specification with arguments
+    input = quote 
+        function some_constraints(x, y)
+            q(x, y) = q(x)q(y)
+            q(x)::PointMass
+        end
+    end
+    output = quote 
+        function some_constraints(x, y;)
+            __constraints__ = GraphPPL.Constraints()
+            q(x, y) = q(x)q(y)
+            q(x)::PointMass
+            return __constraints__
+        end
+    end
+    @test_expression_generating add_constraints_construction(input) output
+
+    # Test 5: add_constraints_construction to constraint specification with function specification with arguments and kwargs
+    input = quote 
+        function some_constraints(x, y; z=1)
+            q(x, y) = q(x)q(y)
+            q(x)::PointMass
+        end
+    end
+    output = quote 
+        function some_constraints(x, y; z=1)
+            __constraints__ = GraphPPL.Constraints()
+            q(x, y) = q(x)q(y)
+            q(x)::PointMass
+            return __constraints__
+        end
+    end
+    @test_expression_generating add_constraints_construction(input) output
 end
 
 @testitem "replace_begin_end" begin