From a1fdab0897982b4edbaef5fcd1be5a2e6f80b7ec Mon Sep 17 00:00:00 2001
From: Harmen Stoppels <harmenstoppels@gmail.com>
Date: Sat, 9 Jun 2018 23:58:21 +0300
Subject: [PATCH] Add test macro for relaxed inference of small unions

---
 stdlib/Test/src/Test.jl | 87 ++++++++++++++++++++++++++++++-----------
 1 file changed, 65 insertions(+), 22 deletions(-)

diff --git a/stdlib/Test/src/Test.jl b/stdlib/Test/src/Test.jl
index 40c0d58ff1184..3a3ed97011602 100644
--- a/stdlib/Test/src/Test.jl
+++ b/stdlib/Test/src/Test.jl
@@ -21,7 +21,7 @@ export @test, @test_throws, @test_broken, @test_skip,
 
 export @testset
 # Legacy approximate testing functions, yet to be included
-export @inferred
+export @inferred, @weakly_inferred
 export detect_ambiguities, detect_unbound_args
 export GenericString, GenericSet, GenericDict, GenericArray
 export guardsrand, TestSetException
@@ -30,6 +30,7 @@ import Distributed: myid
 
 using Random: srand, AbstractRNG, GLOBAL_RNG
 using InteractiveUtils: gen_call_with_extracted_types
+using Core.Compiler: typesubtract
 
 #-----------------------------------------------------------------------
 
@@ -1257,6 +1258,29 @@ function get_testset_depth()
 end
 
 _args_and_call(args...; kwargs...) = (args[1:end-1], kwargs, args[end](args[1:end-1]...; kwargs...))
+
+function call_function_and_infer_type(ex, context)
+    if Meta.isexpr(ex, :ref)
+        ex = Expr(:call, :getindex, ex.args...)
+    end
+    Meta.isexpr(ex, :call) || error("inference test requires a call expression")
+    if any(a->(Meta.isexpr(a, :kw) || Meta.isexpr(a, :parameters)), ex.args)
+        # Has keywords
+        args, kwargs = gensym(), gensym()
+        quote
+            $(esc(args)), $(esc(kwargs)), result = $(esc(Expr(:call, _args_and_call, ex.args[2:end]..., ex.args[1])))
+            inftypes = $(gen_call_with_extracted_types(context, Base.return_types, :($(ex.args[1])($(args)...; $(kwargs)...))))
+        end
+    else
+        # No keywords
+        quote
+            args = ($([esc(ex.args[i]) for i = 2:length(ex.args)]...),)
+            result = $(esc(ex.args[1]))(args...)
+            inftypes = Base.return_types($(esc(ex.args[1])), Base.typesof(args...))
+        end
+    end
+end
+
 """
     @inferred f(x)
 
@@ -1291,29 +1315,9 @@ julia> @inferred max(1,2)
 ```
 """
 macro inferred(ex)
-    if Meta.isexpr(ex, :ref)
-        ex = Expr(:call, :getindex, ex.args...)
-    end
-    Meta.isexpr(ex, :call)|| error("@inferred requires a call expression")
-
     Base.remove_linenums!(quote
         let
-            $(if any(a->(Meta.isexpr(a, :kw) || Meta.isexpr(a, :parameters)), ex.args)
-                # Has keywords
-                args = gensym()
-                kwargs = gensym()
-                quote
-                    $(esc(args)), $(esc(kwargs)), result = $(esc(Expr(:call, _args_and_call, ex.args[2:end]..., ex.args[1])))
-                    inftypes = $(gen_call_with_extracted_types(__module__, Base.return_types, :($(ex.args[1])($(args)...; $(kwargs)...))))
-                end
-            else
-                # No keywords
-                quote
-                    args = ($([esc(ex.args[i]) for i = 2:length(ex.args)]...),)
-                    result = $(esc(ex.args[1]))(args...)
-                    inftypes = Base.return_types($(esc(ex.args[1])), Base.typesof(args...))
-                end
-            end)
+            $(call_function_and_infer_type(ex, __module__))
             @assert length(inftypes) == 1
             rettype = isa(result, Type) ? Type{result} : typeof(result)
             rettype == inftypes[1] || error("return type $rettype does not match inferred return type $(inftypes[1])")
@@ -1322,6 +1326,45 @@ macro inferred(ex)
     end)
 end
 
+"""
+    @weakly_inferred f(x)
+
+Relaxes [`@inferred`](@ref) for functions that return small unions in two ways:
+
+1. When `f(x)` is `missing` or `nothing`, the test passes if `Missing` and `Nothing` are
+respectively subtypes of the inferred type;
+2. Otherwise the test passes provided that the value type equals the inferred type modulo
+`Nothing` and `Missing`.
+
+```jldoctest; setup = :(using InteractiveUtils), filter = r"begin\\n(.|\\n)*end"
+julia> produce(a) = a < 18 ? nothing : a
+julia> @inferred produce(10)
+ERROR: return type Nothing does not match inferred return type Union{Nothing, Int64}
+Stacktrace:
+[...]
+
+julia> @weakly_inferred produce(10)
+
+julia> @weakly_inferred produce(20)
+20
+```
+"""
+macro weakly_inferred(ex)
+    Base.remove_linenums!(quote
+        let
+            $(call_function_and_infer_type(ex, __module__))
+            @assert length(inftypes) == 1
+            inftype = inftypes[1]
+            rettype = isa(result, Type) ? Type{result} : typeof(result)
+            inferred = rettype == Nothing && (Nothing <: inftype) ||
+                       rettype == Missing && (Missing <: inftype) ||
+                       rettype == typesubtract(inftype, Union{Nothing, Missing})
+            inferred || error("return type $rettype does not weakly match inferred return type $(inftype)")
+            result
+        end
+    end)
+end
+
 """
     detect_ambiguities(mod1, mod2...; imported=false, recursive=false, ambiguous_bottom=false)