add support for two-arg @maybe_inferred (#171)

Also adds some direct tests for `@maybe_inferred`.

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

Author: simeonschaub

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.7.5"
+version = "0.7.6"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/testers.jl

@@ -247,14 +247,14 @@ function _ensure_not_running_on_functor(f, name)
 end
 
 """
-    @maybe_inferred f(...)
+    @maybe_inferred [Type] f(...)
 
 Like `@inferred`, but does not check the return type if tests are run as part of PkgEval or
 if the environment variable `CHAINRULES_TEST_INFERRED` is set to `false`.
 """
-macro maybe_inferred(ex)
-    inferred = Expr(:macrocall, GlobalRef(Test, Symbol("@inferred")), __source__, ex)
-    return :(TEST_INFERRED[] ? $(esc(inferred)) : $(esc(ex)))
+macro maybe_inferred(ex...)
+    inferred = Expr(:macrocall, GlobalRef(Test, Symbol("@inferred")), __source__, ex...)
+    return :(TEST_INFERRED[] ? $(esc(inferred)) : $(esc(last(ex))))
 end
 
 """

test/testers.jl

@@ -546,4 +546,22 @@ end
         test_frule(f_notimplemented, randn(), randn())
         test_rrule(f_notimplemented, randn(), randn())
     end
+
+    @testset "@maybe_inferred" begin
+        f_noninferrable(x) = Ref{Real}(x)[]
+
+        @test @maybe_inferred(identity(1)) == 1
+        @test errors("return type $Int does not match inferred return type Real") do
+            @maybe_inferred f_noninferrable(1)
+        end
+        @test @maybe_inferred(Real, f_noninferrable(1)) == 1
+
+        ChainRulesTestUtils.TEST_INFERRED[] = false
+
+        @test @maybe_inferred(identity(1)) == 1
+        @test @maybe_inferred(f_noninferrable(1)) == 1
+        @test @maybe_inferred(Real, f_noninferrable(1)) == 1
+
+        ChainRulesTestUtils.TEST_INFERRED[] = true
+    end
 end