Show evaluated test arguments from broadcast functions

Author: omus

stdlib/Test/src/Test.jl

@@ -340,46 +340,55 @@ struct Threw <: ExecutionResult
     source::LineNumberNode
 end
 
-function eval_test(evaluated::Expr, quoted::Expr, source::LineNumberNode, negate::Bool=false)
-    evaled_args = evaluated.args
+function eval_test_comparison(comparisons::Tuple, quoted::Expr, source::LineNumberNode, negate::Bool=false)
     quoted_args = quoted.args
-    n = length(evaled_args)
+    n = length(comparisons)
     kw_suffix = ""
-    if evaluated.head === :comparison
-        args = evaled_args
-        res = true
-        i = 1
-        while i < n
-            a, op, b = args[i], args[i+1], args[i+2]
-            if res
-                res = op(a, b)
-            end
-            quoted_args[i] = a
-            quoted_args[i+2] = b
-            i += 2
-        end
 
-    elseif evaluated.head === :call
-        op = evaled_args[1]
-        kwargs = (evaled_args[2]::Expr).args  # Keyword arguments from `Expr(:parameters, ...)`
-        args = evaled_args[3:n]
-
-        res = op(args...; kwargs...)
-
-        # Create "Evaluated" expression which looks like the original call but has all of
-        # the arguments evaluated
-        func_sym = quoted_args[1]::Union{Symbol,Expr}
-        if isempty(kwargs)
-            quoted = Expr(:call, func_sym, args...)
-        elseif func_sym === :≈ && !res
-            quoted = Expr(:call, func_sym, args...)
-            kw_suffix = " ($(join(["$k=$v" for (k, v) in kwargs], ", ")))"
-        else
-            kwargs_expr = Expr(:parameters, [Expr(:kw, k, v) for (k, v) in kwargs]...)
-            quoted = Expr(:call, func_sym, kwargs_expr, args...)
+    res = true
+    i = 1
+    while i < n
+        a, op, b = comparisons[i], comparisons[i+1], comparisons[i+2]
+        if res
+            res = op(a, b)
         end
+        quoted_args[i] = a
+        quoted_args[i+2] = b
+        i += 2
+    end
+
+    if negate
+        res = !res
+        quoted = Expr(:call, :!, quoted)
+    end
+
+    Returned(res,
+             # stringify arguments in case of failure, for easy remote printing
+             res === true ? quoted : sprint(print, quoted, context=(:limit => true)) * kw_suffix,
+             source)
+end
+
+function eval_test_function(func, args, kwargs, quoted_func::Union{Expr,Symbol}, source::LineNumberNode, negate::Bool=false)
+    res = func(args...; kwargs...)
+
+    # Create "Evaluated" expression which looks like the original call but has all of
+    # the arguments evaluated
+    kw_suffix = ""
+    if quoted_func === :≈ && !res
+        kw_suffix = " ($(join(["$k=$v" for (k, v) in kwargs], ", ")))"
+        quoted_args = args
+    elseif isempty(kwargs)
+        quoted_args = args
     else
-        throw(ArgumentError("Unhandled expression type: $(evaluated.head)"))
+        kwargs_expr = Expr(:parameters, [Expr(:kw, k, v) for (k, v) in kwargs]...)
+        quoted_args = [kwargs_expr, args...]
+    end
+
+    # Properly render broadcast function call syntax, e.g. `(==).(1, 2)` or `Base.:(==).(1, 2)`.
+    quoted = if isa(quoted_func, Expr) && quoted_func.head === :. && length(quoted_func.args) == 1
+        Expr(:., quoted_func.args[1], Expr(:tuple, quoted_args...))
+    else
+        Expr(:call, quoted_func, quoted_args...)
     end
 
     if negate
@@ -576,14 +585,90 @@ macro test_skip(ex, kws...)
     return :(record(get_testset(), $testres))
 end
 
-function _can_escape_call(@nospecialize ex)
-    ex.head === :call || return false
+function _should_escape_call(@nospecialize ex)
+    isa(ex, Expr) || return false
+
+    args = if ex.head === :call
+        ex.args[2:end]
+    elseif ex.head === :. && length(ex.args) == 2 && isa(ex.args[2], Expr) && ex.args[2].head === :tuple
+        # Support for broadcasted function calls (e.g. `(==).(1, 2)`)
+        ex.args[2].args
+    else
+        # Expression is not a function call
+        return false
+    end
+
+    # Avoid further processing on calls without any arguments
+    return length(args) > 0
+end
+
+# Escapes all of the positional arguments and keywords of a function such that we can call
+# the function at runtime.
+function _escape_call(@nospecialize ex)
+    if isa(ex, Expr) && ex.head === :call
+        # Update broadcast comparison calls to the function call syntax
+        # (e.g. `1 .== 1` becomes `(==).(1, 1)`)
+        func_str = string(ex.args[1])
+        escaped_func = if first(func_str) == '.'
+            esc(Expr(:., Symbol(func_str[2:end])))
+        else
+            esc(ex.args[1])
+        end
+        quoted_func = QuoteNode(ex.args[1])
+        args = ex.args[2:end]
+    elseif isa(ex, Expr) && ex.head === :. && length(ex.args) == 2 && isa(ex.args[2], Expr) && ex.args[2].head === :tuple
+        # Support for broadcasted function calls (e.g. `(==).(1, 2)`)
+        escaped_func = if isa(ex.args[1], Expr) && ex.args[1].head == :.
+            Expr(:call, Expr(:., :Broadcast, QuoteNode(:BroadcastFunction)), esc(ex.args[1]))
+        else
+            Expr(:., esc(ex.args[1]))
+        end
+        quoted_func = QuoteNode(Expr(:., ex.args[1]))
+        args = ex.args[2].args
+    else
+        throw(ArgumentError("$ex is not a call expression"))
+    end
+
+    escaped_args = []
+    escaped_kwargs = []
 
-    # Broadcasted functions are not currently supported
-    first(string(ex.args[1])) != '.' || return false
+    # Positional arguments and keywords that occur before `;`. Note that the keywords are
+    # being revised into a form we can splat.
+    for a in args
+        if isa(a, Expr) && a.head === :parameters
+            continue
+        elseif isa(a, Expr) && a.head === :kw
+            # Keywords that occur before `;`. Note that the keywords are being revised into
+            # a form we can splat.
+            push!(escaped_kwargs, Expr(:call, :(=>), QuoteNode(a.args[1]), esc(a.args[2])))
+        elseif isa(a, Expr) && a.head === :...
+            push!(escaped_args, Expr(:..., esc(a.args[1])))
+        else
+            push!(escaped_args, esc(a))
+        end
+    end
 
-    # At least one positional argument or keyword
-    return length(ex.args) > 1
+    # Keywords that occur after ';'
+    if length(args) > 0 && isa(args[1], Expr) && args[1].head === :parameters
+        for kw in args[1].args
+            if isa(kw, Expr) && kw.head === :kw
+                push!(escaped_kwargs, Expr(:call, :(=>), QuoteNode(kw.args[1]), esc(kw.args[2])))
+            elseif isa(kw, Expr) && kw.head === :...
+                push!(escaped_kwargs, Expr(:..., esc(kw.args[1])))
+            elseif isa(kw, Expr) && kw.head === :.
+                push!(escaped_kwargs, Expr(:call, :(=>), QuoteNode(kw.args[2].value), esc(Expr(:., kw.args[1], QuoteNode(kw.args[2].value)))))
+            elseif isa(kw, Symbol)
+                push!(escaped_kwargs, Expr(:call, :(=>), QuoteNode(kw), esc(kw)))
+            end
+        end
+    end
+
+    return (;
+        func=escaped_func,
+        args=escaped_args,
+        kwargs=escaped_kwargs,
+        quoted_func,
+    )
 end
 
 # An internal function, called by the code generated by the @test
@@ -613,60 +698,22 @@ function get_test_result(ex, source)
         ex = Expr(:comparison, ex.args[1], ex.head, ex.args[2])
     end
     if isa(ex, Expr) && ex.head === :comparison
-        # pass all terms of the comparison to `eval_comparison`, as an Expr
+        # pass all terms of the comparison to `eval_test_comparison`, as a tuple
         escaped_terms = [esc(arg) for arg in ex.args]
         quoted_terms = [QuoteNode(arg) for arg in ex.args]
-        testret = :(eval_test(
-            Expr(:comparison, $(escaped_terms...)),
+        testret = :(eval_test_comparison(
+            ($(escaped_terms...),),
             Expr(:comparison, $(quoted_terms...)),
             $(QuoteNode(source)),
             $negate,
         ))
-    elseif isa(ex, Expr) && _can_escape_call(ex)
-        escaped_func = esc(ex.args[1])
-        quoted_func = QuoteNode(ex.args[1])
-
-        escaped_args = []
-        escaped_kwargs = []
-
-        # Keywords that occur before `;`. Note that the keywords are being revised into
-        # a form we can splat.
-        for a in ex.args[2:end]
-            if isa(a, Expr) && a.head === :kw
-                push!(escaped_kwargs, Expr(:call, :(=>), QuoteNode(a.args[1]), esc(a.args[2])))
-            end
-        end
-
-        # Keywords that occur after ';'
-        parameters_expr = ex.args[2]
-        if isa(parameters_expr, Expr) && parameters_expr.head === :parameters
-            for a in parameters_expr.args
-                if isa(a, Expr) && a.head === :kw
-                    push!(escaped_kwargs, Expr(:call, :(=>), QuoteNode(a.args[1]), esc(a.args[2])))
-                elseif isa(a, Expr) && a.head === :...
-                    push!(escaped_kwargs, Expr(:..., esc(a.args[1])))
-                elseif isa(a, Expr) && a.head === :.
-                    push!(escaped_kwargs, Expr(:call, :(=>), QuoteNode(a.args[2].value), esc(Expr(:., a.args[1], QuoteNode(a.args[2].value)))))
-                elseif isa(a, Symbol)
-                    push!(escaped_kwargs, Expr(:call, :(=>), QuoteNode(a), esc(a)))
-                end
-            end
-        end
-
-        # Positional arguments
-        for a in ex.args[2:end]
-            isa(a, Expr) && a.head in (:kw, :parameters) && continue
-
-            if isa(a, Expr) && a.head === :...
-                push!(escaped_args, Expr(:..., esc(a.args[1])))
-            else
-                push!(escaped_args, esc(a))
-            end
-        end
-
-        testret = :(eval_test(
-            Expr(:call, $escaped_func, Expr(:parameters, $(escaped_kwargs...)), $(escaped_args...)),
-            Expr(:call, $quoted_func),
+    elseif _should_escape_call(ex)
+        call = _escape_call(ex)
+        testret = :(eval_test_function(
+            $(call.func),
+            ($(call.args...),),
+            ($(call.kwargs...),),
+            $(call.quoted_func),
             $(QuoteNode(source)),
             $negate,
         ))