Avoid modifying closed module globals in computations documentation.

tanaylab · Apr 2, 2024 · 460032a · 460032a
1 parent 7580d2d
commit 460032a
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 21 deletions.
diff --git a/docs/v0.1.0/.documenter-siteinfo.json b/docs/v0.1.0/.documenter-siteinfo.json
@@ -1 +1 @@
-{"documenter":{"julia_version":"1.10.2","generation_timestamp":"2024-04-01T20:12:39","documenter_version":"1.3.0"}}
+{"documenter":{"julia_version":"1.10.2","generation_timestamp":"2024-04-02T11:52:43","documenter_version":"1.3.0"}}
diff --git a/src/computations.jl b/src/computations.jl
@@ -62,10 +62,15 @@ struct FunctionMetadata
     defaults::Dict{AbstractString, Any}
 end
 
-const METADATA_OF_FUNCTION = Dict{AbstractString, FunctionMetadata}()
-
-function set_metadata_of_function(full_name::AbstractString, function_metadata::FunctionMetadata)::Nothing
-    METADATA_OF_FUNCTION[full_name] = function_metadata
+function set_metadata_of_function(
+    function_module::Module,
+    function_name::Symbol,
+    function_metadata::FunctionMetadata,
+)::Nothing
+    if !isdefined(function_module, :__DAF_FUNCTION_METADATA__)
+        function_module.__DAF_FUNCTION_METADATA__ = Dict{Symbol, FunctionMetadata}()
+    end
+    function_module.__DAF_FUNCTION_METADATA__[function_name] = function_metadata
     return nothing
 end
 
@@ -112,9 +117,14 @@ macro computation(definition)
         error("@logged requires a named function")
     end
     @assert function_name isa Symbol
+    function_module = __module__
+    full_name = "$(function_module).$(function_name)"
 
-    full_name = "$(__module__).$(function_name)"
-    set_metadata_of_function(full_name, FunctionMetadata(Contract[], collect_defaults(inner_definition)))
+    set_metadata_of_function(
+        function_module,
+        function_name,
+        FunctionMetadata(Contract[], collect_defaults(inner_definition)),
+    )
 
     inner_definition[:name] = Symbol(function_name, :_inner)
     outer_definition[:body] = Expr(
@@ -136,15 +146,20 @@ macro computation(contract, definition)
         error("@computation requires a named function")
     end
     @assert function_name isa Symbol
+    function_module = __module__
+    full_name = "$(function_module).$(function_name)"
 
-    full_name = "$(__module__).$(function_name)"
-    set_metadata_of_function(full_name, FunctionMetadata([eval(contract)], collect_defaults(inner_definition)))
+    set_metadata_of_function(
+        function_module,
+        function_name,
+        FunctionMetadata([eval(contract)], collect_defaults(inner_definition)),
+    )
 
     inner_definition[:name] = Symbol(function_name, :_inner)
     outer_definition[:body] = Expr(
         :call,
         :(Daf.Computations.computation_wrapper(
-            Daf.Computations.METADATA_OF_FUNCTION[$full_name].contracts[1],
+            $function_module.__DAF_FUNCTION_METADATA__[Symbol($function_name)].contracts[1],
             $full_name,
             $(ExprTools.combinedef(inner_definition)),
         )),
@@ -164,18 +179,21 @@ macro computation(first_contract, second_contract, definition)
         error("@computation requires a named function")
     end
     @assert function_name isa Symbol
+    function_module = __module__
+    full_name = "$(function_module).$(function_name)"
 
-    full_name = "$(__module__).$(function_name)"
-
-    METADATA_OF_FUNCTION[full_name] =
-        FunctionMetadata([eval(first_contract), eval(second_contract)], collect_defaults(inner_definition))
+    set_metadata_of_function(
+        function_module,
+        function_name,
+        FunctionMetadata([eval(first_contract), eval(second_contract)], collect_defaults(inner_definition)),
+    )
 
     inner_definition[:name] = Symbol(function_name, :_inner)
     outer_definition[:body] = Expr(
         :call,
         :(Daf.Computations.computation_wrapper(
-            Daf.Computations.METADATA_OF_FUNCTION[$full_name].contracts[1],
-            Daf.Computations.METADATA_OF_FUNCTION[$full_name].contracts[2],
+            $function_module.__DAF_FUNCTION_METADATA__[Symbol($function_name)].contracts[1],
+            $function_module.__DAF_FUNCTION_METADATA__[Symbol($function_name)].contracts[2],
             $full_name,
             $(ExprTools.combinedef(inner_definition)),
         )),
@@ -348,15 +366,28 @@ const DEFAULT = DefaultContainer()
 function get_metadata(doc_str::Base.Docs.DocStr)::Tuple{AbstractString, FunctionMetadata}
     binding = doc_str.data[:binding]
     object = Docs.resolve(binding)
-    full_name = "$(doc_str.data[:module]).$(Symbol(object))"
-    metadata = get(METADATA_OF_FUNCTION, full_name, nothing)
+    object_module = nothing
+    for method in methods(object)
+        try
+            object_module = method.module
+            if isdefined(object_module, :__DAF_FUNCTION_METADATA__)
+                break
+            end
+        catch  # untested
+        end
+    end
+    if object_module == nothing
+        metadata = nothing  # untested
+    else
+        metadata = get(object_module.__DAF_FUNCTION_METADATA__, Symbol(object), nothing)
+    end
     if metadata == nothing
         error(
-            "no contract(s) associated with: $(full_name)\n" *
-            "use: @computation Contract(...) function $(full_name)(...)",
+            "no contract(s) associated with: $(object_module).$(object)\n" *
+            "use: @computation Contract(...) function $(object_module).$(object)(...)",
         )
     end
-    return full_name, metadata
+    return "$(object_module).$(object)", metadata
 end
 
 end # module