From 730b5c948488b272c131e02bb5c83876a374cfd2 Mon Sep 17 00:00:00 2001
From: Luc Briand <34173752+Keluaa@users.noreply.github.com>
Date: Sat, 15 Jun 2024 11:02:32 +0200
Subject: [PATCH] Allow any object into the macros, making any extension
 possible

---
 docs/src/extensions.md |  8 ++++++++
 src/compare.jl         | 46 ++++++++++++++++++++++++------------------
 test/runtests.jl       | 12 +++++------
 3 files changed, 40 insertions(+), 26 deletions(-)

diff --git a/docs/src/extensions.md b/docs/src/extensions.md
index 0ea236d..9be7b0a 100644
--- a/docs/src/extensions.md
+++ b/docs/src/extensions.md
@@ -57,3 +57,11 @@ involves two functions:
 - `CodeDiffs.extract_extra_options(f, kwargs)` returns some additional `kwargs` which are passed to `get_code`
 - `CodeDiffs.get_code(code_type, f, types; kwargs...)` allows to change `f` depending on its type.
   To avoid method ambiguities, do not put type constraints on `code_type`.
+
+Defining a new object type which can be put as an argument to `@code_diff` or `@code_for`
+invoves at one function: `CodeDiffs.code_for_diff(obj::YourType; kwargs...)`.
+It must return two `String`s, one without and the other without highlighting.
+When calling `@code_for obj`, [`code_for_diff(obj)`](@ref) will be called only if `obj` is
+not a call expression or a quoted `Expr`.
+`kwargs` are the options passed to `@code_for` or the options passed to `@code_diff` for
+the side of `obj`.
diff --git a/src/compare.jl b/src/compare.jl
index 83f8cf7..7d66367 100644
--- a/src/compare.jl
+++ b/src/compare.jl
@@ -100,12 +100,13 @@ argconvert(@nospecialize(code_type), arg) = arg
 extract_extra_options(@nospecialize(f), _) = (;)
 separate_kwargs(code_type::Val, args...; kwargs...) = (argconvert.(Ref(code_type), args), values(kwargs))
 is_error_expr(expr) = Base.isexpr(expr, :call) && expr.args[1] in (:error, :throw)
+is_call_expr(expr) = Base.isexpr(expr, :call) || (Base.isexpr(expr, :(.), 2) && Base.isexpr(expr.args[2], :tuple))
 
 
 function gen_code_for_diff_call(mod, expr, diff_options)
     # `diff_options` must be already `esc`aped, but not `expr`
 
-    if !Base.isexpr(expr, [:call, :(.)])
+    if !is_call_expr(expr)
         error_str = "Expected call (or dot call) to function, got: $expr"
         return :(throw(ArgumentError($error_str)))
     end
@@ -173,6 +174,23 @@ function gen_code_for_diff_call(mod, expr, diff_options)
 end
 
 
+function code_diff_for_expr(code_expr, options, mod)
+    if is_call_expr(code_expr)
+        # Generic call comparison
+        return gen_code_for_diff_call(mod, code_expr, options)
+    else
+        if Base.isexpr(code_expr, :quote)
+            # AST comparison
+            type_guess = (Expr(:kw, :type, QuoteNode(:ast)),)
+        else
+            # Mystery comparison
+            type_guess = ()
+        end
+        return :($code_for_diff($(esc(code_expr)); $(type_guess...), $(options...)))
+    end
+end
+
+
 """
     @code_diff [type=:native] [color=true] [cleanup=true] [option=value...] f₁(...) f₂(...)
     @code_diff [option=value...] :(expr₁) :(expr₂)
@@ -256,17 +274,10 @@ macro code_diff(args...)
     code₁ isa QuoteNode && (code₁ = Expr(:quote, Expr(:block, code₁.value)))
     code₂ isa QuoteNode && (code₂ = Expr(:quote, Expr(:block, code₂.value)))
 
-    if Base.isexpr(code₁, :quote) && Base.isexpr(code₂, :quote)
-        code₁ = esc(code₁)
-        code₂ = esc(code₂)
-        code_for_diff₁ = :($code_for_diff($code₁; type=:ast, $(options₁...)))
-        code_for_diff₂ = :($code_for_diff($code₂; type=:ast, $(options₂...)))
-    else
-        code_for_diff₁ = gen_code_for_diff_call(__module__, code₁, options₁)
-        code_for_diff₂ = gen_code_for_diff_call(__module__, code₂, options₂)
-        is_error_expr(code_for_diff₁) && return code_for_diff₁
-        is_error_expr(code_for_diff₂) && return code_for_diff₂
-    end
+    code_for_diff₁ = code_diff_for_expr(code₁, options₁, __module__)
+    code_for_diff₂ = code_diff_for_expr(code₂, options₂, __module__)
+    is_error_expr(code_for_diff₁) && return code_for_diff₁
+    is_error_expr(code_for_diff₂) && return code_for_diff₂
 
     return quote
         let
@@ -342,15 +353,10 @@ macro code_for(args...)
     # Simple values such as `:(1)` are stored in a `QuoteNode`
     code isa QuoteNode && (code = Expr(:quote, Expr(:block, code.value)))
 
+    # Place the diff options in a temp variable in order to extract `io` at runtime
     diff_options = esc(gensym(:diff_options))
-
-    if Base.isexpr(code, :quote)
-        code = esc(code)
-        code_for = :($code_for_diff($code; type=:ast, $(options...)))
-    else
-        code_for = gen_code_for_diff_call(__module__, code, [:($diff_options...)])
-        is_error_expr(code_for) && return code_for
-    end
+    code_for = code_diff_for_expr(code, [:($diff_options...)], __module__)
+    is_error_expr(code_for) && return code_for
 
     return quote
         let
diff --git a/test/runtests.jl b/test/runtests.jl
index 08583b6..dc6c021 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -461,15 +461,15 @@ end
 
     @testset "Macros" begin
         @testset "error" begin
-            @test_throws "Expected call" @code_diff "f()" g()
-            @test_throws "Expected call" @code_diff f() "g()"
-            @test_throws "Expected call" @code_diff "f()" "g()"
-            @test_throws "Expected call" @code_diff a b
+            @test_throws MethodError @code_diff "f()" g()
+            @test_throws MethodError @code_diff f() "g()"
+            @test_throws MethodError @code_diff "f()" "g()"
+            @test_throws UndefVarError @code_diff a b
             @test_throws "`key=value`, got: `a + 1`" @code_diff a+1 b c
             @test_throws "world age" @code_diff type=:ast world_1=1 f() f()
 
-            @test_throws "Expected call" @code_for "f()"
-            @test_throws "Expected call" @code_for a
+            @test_throws MethodError @code_for "f()"
+            @test_throws UndefVarError @code_for a
             @test_throws "`key=value`, got: `a + 1`" @code_for a+1 b c
             @test_throws "world age" @code_for type=:ast world=1 f()
         end