test: include test of ambiguous operator error

Author: MilesCranmer

src/OperatorEnumConstruction.jl

@@ -118,6 +118,20 @@ function _unpack_broadcast_function(f)
     end
 end
 
+function _validate_no_ambiguous_broadcasts(operators::AbstractOperatorEnum)
+    for ops in (operators.binops, operators.unaops), op in ops
+        if op isa Broadcast.BroadcastFunction &&
+            (op.f in operators.binops || op.f in operators.unaops)
+            throw(
+                ArgumentError(
+                    "Usage of both broadcasted and unbroadcasted operator `$(op.f)` is ambiguous",
+                ),
+            )
+        end
+    end
+    return nothing
+end
+
 function empty_all_globals!(; force=true)
     if force || islocked(LATEST_LOCK)
         lock(LATEST_LOCK) do
@@ -280,6 +294,7 @@ function _extend_operators(operators, skip_user_operators, kws, __module__::Modu
     #! format: off
     return quote
         local $type_requirements, $build_converters, $binary_exists, $unary_exists
+        $(_validate_no_ambiguous_broadcasts)($operators)
         lock($LATEST_LOCK)
         if isa($operators, $OperatorEnum)
             $type_requirements = $(on_type == nothing ? Number : on_type)
@@ -437,17 +452,6 @@ redefine operators for `AbstractExpressionNode` types, as well as `show`, `print
         end
     end
 
-    for ops in (binary_operators, unary_operators), op in ops
-        if op isa Broadcast.BroadcastFunction &&
-            (op.f in binary_operators || op.f in unary_operators)
-            throw(
-                ArgumentError(
-                    "Usage of both broadcasted and unbroadcasted operator $(op.f) is ambiguous",
-                ),
-            )
-        end
-    end
-
     operators = OperatorEnum(Tuple(binary_operators), Tuple(unary_operators))
 
     if define_helper_functions

test/test_broadcasted_operators.jl

@@ -0,0 +1,37 @@
+@testitem "Test broadcasted operators in enum" begin
+    using DynamicExpressions
+
+    operators = OperatorEnum(;
+        binary_operators=[*, +, (.-)],
+        unary_operators=[sin, Broadcast.BroadcastFunction(cos)],
+    )
+    @extend_operators operators
+
+    x1 = Node{Float64}(; feature=1)
+    x2 = Node{Float64}(; feature=2)
+
+    @test x1 + x2 isa Node{Float64}
+    @test repr(x1 + x2) == "x1 + x2"
+    @test x1 - x2 isa Node{Float64}
+    @test repr(x1 - x2) == "x1 .- x2"
+
+    @test cos(x1) isa Node{Float64}
+    @test repr(cos(x1)) == "cos.(x1)"
+end
+
+@testitem "Test error upon ambiguous operators" begin
+    using DynamicExpressions
+
+    operators = OperatorEnum(;
+        unary_operators=[cos, Broadcast.BroadcastFunction(cos)],
+        define_helper_functions=false,
+    )
+    @test_throws ArgumentError @eval begin
+        @extend_operators operators
+    end
+    if VERSION >= v"1.9"
+        @test_throws "Usage of both broadcasted and unbroadcasted operator `cos` is ambiguous" @eval begin
+            @extend_operators operators
+        end
+    end
+end

test/unittest.jl

@@ -62,6 +62,8 @@ end
     include("test_non_number_eval_tree_array.jl")
 end
 
+include("test_broadcasted_operators.jl")
+
 @testitem "Test hash of tree" begin
     include("test_hash.jl")
 end