Merge pull request #182 from ReactiveBayes/dev-4.0.0-get-extra-or-def…

…ault Add getextra with default
ReactiveBayes · Mar 19, 2024 · 37c54e3 · 37c54e3
2 parents f1c28b4 + 7835eba
commit 37c54e3
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/src/graph_engine.jl b/src/graph_engine.jl
@@ -549,6 +549,7 @@ getextra(node::NodeData)      = node.extra
 
 hasextra(node::NodeData, key::Symbol) = haskey(node.extra, key)
 getextra(node::NodeData, key::Symbol) = getindex(node.extra, key)
+getextra(node::NodeData, key::Symbol, default) = hasextra(node, key) ? getextra(node, key) : default
 setextra!(node::NodeData, key::Symbol, value) = insert!(node.extra, key, value)
 
 """
@@ -562,6 +563,9 @@ end
 function getextra(node::NodeData, key::NodeDataExtraKey{K, T})::T where {K, T}
     return getindex(node.extra, K)::T
 end
+function getextra(node::NodeData, key::NodeDataExtraKey{K, T}, default::T)::T where {K, T}
+    return hasextra(node, key) ? (getextra(node, key)::T) : default
+end
 function setextra!(node::NodeData, key::NodeDataExtraKey{K}, value::T) where {K, T}
     return insert!(node.extra, K, value)
 end

diff --git a/test/graph_engine_tests.jl b/test/graph_engine_tests.jl
@@ -127,9 +127,15 @@ end
         nodedata = NodeData(context, properties)
 
         @test !hasextra(nodedata, :a)
+        @test getextra(nodedata, :a, 2) === 2
+        @test !hasextra(nodedata, :a) # the default should not add the extra property, only return
         setextra!(nodedata, :a, 1)
         @test hasextra(nodedata, :a)
         @test getextra(nodedata, :a) === 1
+        @test getextra(nodedata, :a, 2) === 1
+        @test !hasextra(nodedata, :b)
+        @test_throws Exception getextra(nodedata, :b)
+        @test getextra(nodedata, :b, 2) === 2
 
         # In the current implementation it is not possible to update extra properties
         @test_throws Exception setextra!(nodedata, :a, 2)
@@ -142,9 +148,14 @@ end
         constkey_c_float = NodeDataExtraKey{:c, Float64}()
 
         @test !@inferred(hasextra(nodedata, constkey_c_float))
+        @test @inferred(getextra(nodedata, constkey_c_float, 4.0)) === 4.0
         @inferred(setextra!(nodedata, constkey_c_float, 3.0))
         @test @inferred(hasextra(nodedata, constkey_c_float))
         @test @inferred(getextra(nodedata, constkey_c_float)) === 3.0
+        @test @inferred(getextra(nodedata, constkey_c_float, 4.0)) === 3.0
+
+        # The default has a different type from the key (4.0 is Float and 4 is Int), thus the error 
+        @test_throws MethodError getextra(nodedata, constkey_c_float, 4)
 
         constkey_d_int = NodeDataExtraKey{:d, Int64}()