fix: preserve symbolic external units across precompile

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

Author: MilesCranmerBot

src/symbolic_dimensions.jl

@@ -426,13 +426,26 @@ module SymbolicUnits
 
     import ..UNIT_SYMBOLS
     import ..CONSTANT_SYMBOLS
+    import ..ALL_MAPPING
     import ..SymbolicDimensionsSingleton
     import ..constructorof
     import ..DEFAULT_SYMBOLIC_QUANTITY_TYPE
     import ..DEFAULT_SYMBOLIC_QUANTITY_OUTPUT_TYPE
     import ..DEFAULT_VALUE_TYPE
     import ..DEFAULT_DIM_BASE_TYPE
+    import ..INDEX_TYPE
+    import ..UnionAbstractQuantity
     import ..WriteOnceReadMany
+    import ..disambiguate_constant_symbol
+
+    symbolic_unit_from_symbol(unit::Symbol) = constructorof(DEFAULT_SYMBOLIC_QUANTITY_TYPE)(
+        DEFAULT_VALUE_TYPE(1.0),
+        SymbolicDimensionsSingleton{DEFAULT_DIM_BASE_TYPE}(unit)
+    )
+    symbolic_constant_from_symbol(unit::Symbol) = constructorof(DEFAULT_SYMBOLIC_QUANTITY_TYPE)(
+        DEFAULT_VALUE_TYPE(1.0),
+        SymbolicDimensionsSingleton{DEFAULT_DIM_BASE_TYPE}(disambiguate_constant_symbol(unit))
+    )
 
     # Lazily create unit symbols (since there are so many)
     module Constants
@@ -479,11 +492,7 @@ module SymbolicUnits
     # Non-eval version of `update_symbolic_unit_values!` for registering units in
     # an external module.
     function update_external_symbolic_unit_value(unit)
-        unit = constructorof(DEFAULT_SYMBOLIC_QUANTITY_TYPE)(
-            DEFAULT_VALUE_TYPE(1.0),
-            SymbolicDimensionsSingleton{DEFAULT_DIM_BASE_TYPE}(unit)
-        )
-        push!(SYMBOLIC_UNIT_VALUES, unit)
+        push!(SYMBOLIC_UNIT_VALUES, symbolic_unit_from_symbol(unit))
     end
 
     """
@@ -512,39 +521,50 @@ module SymbolicUnits
     as_quantity(x::Number) = convert(DEFAULT_SYMBOLIC_QUANTITY_OUTPUT_TYPE, x)
     as_quantity(x) = error("Unexpected type evaluated: $(typeof(x))")
 
-    @unstable function map_to_scope(ex::Expr)
+    @unstable map_to_scope(ex::Expr) = map_to_scope(@__MODULE__, ex)
+    @unstable function map_to_scope(mod::Module, ex::Expr)
         if !(ex.head == :call) && !(ex.head == :. && ex.args[1] == :Constants)
             throw(ArgumentError("Unexpected expression: $ex. Only `:call` and `:.` (for `SymbolicConstants`) are expected."))
         end
         if ex.head == :call
-            ex.args[2:end] = map(map_to_scope, ex.args[2:end])
+            ex.args[2:end] = map(arg -> map_to_scope(mod, arg), ex.args[2:end])
             return ex
         else # if ex.head == :. && ex.args[1] == :Constants
             @assert ex.args[2] isa QuoteNode
-            return lookup_constant(ex.args[2].value)
+            return Expr(:call, GlobalRef(@__MODULE__, :lookup_constant), QuoteNode(ex.args[2].value))
         end
     end
-    function map_to_scope(sym::Symbol)
-        if sym in UNIT_SYMBOLS
+    map_to_scope(sym::Symbol) = map_to_scope(@__MODULE__, sym)
+    function map_to_scope(mod::Module, sym::Symbol)
+        if sym in UNIT_SYMBOLS || _has_quantity_binding(mod, sym)
             # return at end
         elseif sym in CONSTANT_SYMBOLS
             throw(ArgumentError("Symbol $sym found in `Constants` but not `Units`. Please use `us\"Constants.$sym\"` instead."))
         else
             throw(ArgumentError("Symbol $sym not found in `Units` or `Constants`."))
         end
-        return lookup_unit(sym)
+        return Expr(:call, GlobalRef(@__MODULE__, :lookup_unit), mod, QuoteNode(sym))
     end
-    function map_to_scope(ex)
-        return ex
+    map_to_scope(ex) = ex
+    map_to_scope(::Module, ex) = ex
+
+    function _has_quantity_binding(mod::Module, sym::Symbol)
+        return isdefined(mod, sym) && Base.invokelatest(getproperty, mod, sym) isa UnionAbstractQuantity
     end
-    function lookup_unit(ex::Symbol)
-        i = findfirst(==(ex), UNIT_SYMBOLS)::Int
-        return as_quantity(SYMBOLIC_UNIT_VALUES[i])
+    function _ensure_registered(mod::Module, sym::Symbol)
+        if iszero(get(ALL_MAPPING, sym, INDEX_TYPE(0))) && _has_quantity_binding(mod, sym)
+            update_all_values = getfield(parentmodule(@__MODULE__), :update_all_values)
+            unit = Base.invokelatest(getproperty, mod, sym)
+            Base.invokelatest(update_all_values, sym, unit)
+        end
+        return nothing
     end
-    function lookup_constant(ex::Symbol)
-        i = findfirst(==(ex), CONSTANT_SYMBOLS)::Int
-        return as_quantity(SYMBOLIC_CONSTANT_VALUES[i])
+    function lookup_unit(mod::Module, ex::Symbol)
+        _ensure_registered(mod, ex)
+        return as_quantity(symbolic_unit_from_symbol(ex))
     end
+    lookup_unit(ex::Symbol) = lookup_unit(@__MODULE__, ex)
+    lookup_constant(ex::Symbol) = as_quantity(symbolic_constant_from_symbol(ex))
 end
 
 import .SymbolicUnits: as_quantity, sym_uparse, SymbolicConstants, map_to_scope
@@ -565,7 +585,7 @@ module. So, for example, `us"Constants.c^2 * Hz^2"` would evaluate to
 namespace collisions, a few physical constants are automatically converted.
 """
 macro us_str(s)
-    ex = map_to_scope(Meta.parse(s))
+    ex = map_to_scope(__module__, Meta.parse(s))
     ex = :($as_quantity($ex))
     return esc(ex)
 end

test/precompile_test/ExternalUnitRegistration.jl

@@ -7,15 +7,21 @@ using Test
 
 @register_unit MyWb u"m^2*kg*s^-2*A^-1"
 
+expanded_mywb() = 1u"MyWb"
+symbolic_mywb() = 1us"MyWb"
+
 @testset "Register Unit Inside a Module" begin
     for collection in (UNIT_SYMBOLS, ALL_SYMBOLS, keys(ALL_MAPPING._raw_data), keys(UNIT_MAPPING._raw_data))
         @test :MyWb ∈ collection
     end
 
-    w = u"MyWb"
-    ws = us"MyWb"
+    w = expanded_mywb()
+    ws = symbolic_mywb()
     @test w isa DEFAULT_QUANTITY_TYPE
     @test ws isa DEFAULT_SYMBOLIC_QUANTITY_OUTPUT_TYPE
+    @test w == u"MyWb"
+    @test ws == us"MyWb"
+    @test string(ws) == "1.0 MyWb"
 end
 
 end

test/unittests.jl

@@ -2343,10 +2343,13 @@ end
 
 push!(LOAD_PATH, joinpath(@__DIR__, "precompile_test"))
 
-using ExternalUnitRegistration: MyWb
+using ExternalUnitRegistration: MyWb, expanded_mywb, symbolic_mywb
 @testset "Type of External Unit" begin
     @test MyWb isa DEFAULT_QUANTITY_TYPE
     @test MyWb/u"m^2*kg*s^-2*A^-1" == 1.0
+    @test expanded_mywb() == MyWb
+    @test uexpand(symbolic_mywb()) == MyWb
+    @test string(symbolic_mywb()) == "1.0 MyWb"
 end
 
 pop!(LOAD_PATH)