From 7724fd403c6dce607c27eb377d20b624a0d9c763 Mon Sep 17 00:00:00 2001
From: Chris Foster <chris.foster@tri.edu.au>
Date: Wed, 12 Dec 2018 13:27:16 +1000
Subject: [PATCH] Make u_str work with precompilation

We can check that the unit extension module is loaded into the current
module in which the u_str macro is invoked. This should allow `u_str` to
be safely used with precompilation.
---
 src/user.jl      | 39 ++++++++++++++++++++++++++-------------
 test/runtests.jl | 34 ++++++++++++++++++++--------------
 2 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/src/user.jl b/src/user.jl
index 90d9d153..1d271ff4 100644
--- a/src/user.jl
+++ b/src/user.jl
@@ -471,9 +471,6 @@ will be used.
 Note that what goes inside must be parsable as a valid Julia expression.
 In other words, u"N m" will fail if you intended to write u"N*m".
 
-Since `@u_str` dynamically looks through registered modules, it does not play well with
-precompilation. You shouldn't use this macro in a precompiled module.
-
 Examples:
 
 ```jldoctest
@@ -495,25 +492,25 @@ julia> u"ħ"
 """
 macro u_str(unit)
     ex = Meta.parse(unit)
-    esc(replace_value(ex))
+    esc(replace_value(__module__, ex))
 end
 
 const allowed_funcs = [:*, :/, :^, :sqrt, :√, :+, :-, ://]
-function replace_value(ex::Expr)
+function replace_value(targetmod, ex::Expr)
     if ex.head == :call
         ex.args[1] in allowed_funcs ||
             error("""$(ex.args[1]) is not a valid function call when parsing a unit.
              Only the following functions are allowed: $allowed_funcs""")
         for i=2:length(ex.args)
             if typeof(ex.args[i])==Symbol || typeof(ex.args[i])==Expr
-                ex.args[i]=replace_value(ex.args[i])
+                ex.args[i]=replace_value(targetmod, ex.args[i])
             end
         end
         return ex
     elseif ex.head == :tuple
         for i=1:length(ex.args)
             if typeof(ex.args[i])==Symbol
-                ex.args[i]=replace_value(ex.args[i])
+                ex.args[i]=replace_value(targetmod, ex.args[i])
             else
                 error("only use symbols inside the tuple.")
             end
@@ -524,11 +521,27 @@ function replace_value(ex::Expr)
     end
 end
 
-function replace_value(sym::Symbol)
-    f = m->(isdefined(m,sym) && ustrcheck_bool(getfield(m, sym)))
-    inds = findall(f, unitmodules)
-    isempty(inds) &&
-        error("Symbol $sym could not be found in registered unit modules.")
+function replace_value(targetmod, sym::Symbol)
+    f = m->()
+    inds = findall(unitmodules) do m
+        # Ensure that both the unit exists in the registered unit module, and
+        # that the target module (the one invoking `u_str`) has loaded it so
+        # that precompilation will work.
+        isdefined(m,sym)             && ustrcheck_bool(getfield(m, sym))  &&
+        isdefined(targetmod, nameof(m)) && getfield(targetmod,nameof(m)) === m
+    end
+    if isempty(inds)
+        # Check whether unit exists in the global list to give an improved
+        # error message.
+        f = m->(isdefined(m,sym) && ustrcheck_bool(getfield(m, sym)))
+        hintidx = findfirst(f, unitmodules)
+        if hintidx !== nothing
+            hintmod = unitmodules[hintidx]
+            error("Symbol `$sym` was found in unit module $hintmod, but was not loaded into $targetmod. Consider `using $hintmod` within `$targetmod`?")
+        else
+            error("Symbol $sym could not be found in registered unit modules.")
+        end
+    end
 
     m = unitmodules[inds[end]]
     u = getfield(m, sym)
@@ -539,7 +552,7 @@ function replace_value(sym::Symbol)
     return u
 end
 
-replace_value(literal::Number) = literal
+replace_value(targetmod, literal::Number) = literal
 
 ustrcheck_bool(::Number) = true
 ustrcheck_bool(::MixedUnits) = true
diff --git a/test/runtests.jl b/test/runtests.jl
index d1b59bed..8e48c107 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1436,21 +1436,27 @@ module ShadowUnits
     Unitful.register(ShadowUnits)
 end
 
-let fname = tempname()
-    try
-        ret = open(fname, "w") do f
-            redirect_stderr(f) do
-                # wrap in eval to catch the STDERR output...
-                @test eval(:(typeof(u"m"))) == Unitful.FreeUnits{
-                    (Unitful.Unit{:MyMeter, 𝐋}(0, 1//1),), 𝐋, nothing}
-            end
-        end
-    finally
-        rm(fname, force=true)
-    end
-end
+@test (@test_logs (:warn, r"found in multiple") eval(:(typeof(u"m")))) ==
+    Unitful.FreeUnits{(Unitful.Unit{:MyMeter, 𝐋}(0, 1//1),), 𝐋, nothing}
 
-@test_logs (:warn, r"found in multiple") eval(:(u"m"))
+# Test that the @u_str macro will not find units in modules which are
+# not loaded before the u_str invocation.
+module FooUnits
+    using Unitful
+    @unit foo "foo" MyFoo 1u"m" false
+    Unitful.register(FooUnits)
+end
+# NB: The following is LoadError in 1.0 but an ErrorException in 0.7.
+@test_throws Exception eval(:(module ShouldUseFooUnits
+                                  using Unitful
+                                  foo() = 1u"foo"
+                              end))
+# Test that u_str works when FooUnits is correctly loaded.
+module DoesUseFooUnits
+    using Unitful, ..FooUnits
+    foo() = 1u"foo"
+end
+@test DoesUseFooUnits.foo() === 1u"foo"
 
 # Test to make sure user macros are working properly
 module TUM