feat: allow separate operator names for pretty printing

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -69,7 +69,7 @@ import .NodeModule:
     get_scalar_constants,
     set_scalar_constants!
 @reexport import .StringsModule: string_tree, print_tree
-import .StringsModule: get_op_name
+import .StringsModule: get_op_name, get_pretty_op_name
 @reexport import .OperatorEnumModule: AbstractOperatorEnum
 @reexport import .OperatorEnumConstructionModule:
     OperatorEnum, GenericOperatorEnum, @extend_operators, set_default_variable_names!

src/Strings.jl

@@ -4,19 +4,24 @@ using ..UtilsModule: deprecate_varmap
 using ..OperatorEnumModule: AbstractOperatorEnum
 using ..NodeModule: AbstractExpressionNode, tree_mapreduce
 
-function dispatch_op_name(::Val{deg}, ::Nothing, idx)::Vector{Char} where {deg}
-    if deg == 1
-        return vcat(collect("unary_operator["), collect(string(idx)), [']'])
-    else
-        return vcat(collect("binary_operator["), collect(string(idx)), [']'])
-    end
+function dispatch_op_name(
+    ::Val{deg}, ::Nothing, idx, pretty::Bool
+)::Vector{Char} where {deg}
+    return vcat(
+        collect(deg == 1 ? "unary_operator[" : "binary_operator["),
+        collect(string(idx)),
+        [']'],
+    )
 end
-function dispatch_op_name(::Val{deg}, operators::AbstractOperatorEnum, idx) where {deg}
-    if deg == 1
-        return collect(get_op_name(operators.unaops[idx])::String)
+function dispatch_op_name(
+    ::Val{deg}, operators::AbstractOperatorEnum, idx, pretty::Bool
+) where {deg}
+    op = if deg == 1
+        operators.unaops[idx]
     else
-        return collect(get_op_name(operators.binops[idx])::String)
+        operators.binops[idx]
     end
+    return collect((pretty ? get_pretty_op_name(op) : get_op_name(op))::String)
 end
 
 const OP_NAME_CACHE = (; x=Dict{UInt64,String}(), lock=Threads.SpinLock())
@@ -47,6 +52,9 @@ function get_op_name(op::F) where {F}
         unlock(OP_NAME_CACHE.lock)
     end
 end
+function get_pretty_op_name(op::F) where {F}
+    return get_op_name(op)
+end
 
 @inline function strip_brackets(s::Vector{Char})::Vector{Char}
     if first(s) == '(' && last(s) == ')'
@@ -145,8 +153,9 @@ function string_tree(
     raw::Union{Bool,Nothing}=nothing,
     varMap=nothing,
 )::String where {T,F1<:Function,F2<:Function}
-    !isnothing(raw) &&
+    if !isnothing(raw)
         Base.depwarn("`raw` is deprecated; use `pretty` instead", :string_tree)
+    end
     pretty = @something(pretty, _not(raw), false)
     variable_names = deprecate_varmap(variable_names, varMap, :string_tree)
     raw_output = tree_mapreduce(
@@ -162,9 +171,9 @@ function string_tree(
         end,
         let operators = operators
             (branch,) -> if branch.degree == 1
-                dispatch_op_name(Val(1), operators, branch.op)::Vector{Char}
+                dispatch_op_name(Val(1), operators, branch.op, pretty)::Vector{Char}
             else
-                dispatch_op_name(Val(2), operators, branch.op)::Vector{Char}
+                dispatch_op_name(Val(2), operators, branch.op, pretty)::Vector{Char}
             end
         end,
         combine_op_with_inputs,

test/test_print.jl

@@ -117,3 +117,61 @@ end
     @test string(tree) == "(k1 * x2) + x3"
     empty!(DynamicExpressions.OperatorEnumConstructionModule.LATEST_VARIABLE_NAMES.x)
 end
+
+@testset "Test pretty format for operators" begin
+    # Define a custom operator with different pretty representation
+    @eval begin
+        my_pretty_op(x, y) = x + y
+        DE.get_op_name(::typeof(my_pretty_op)) = "my_pretty_op"
+        DE.get_pretty_op_name(::typeof(my_pretty_op)) = "pretty_op_two"
+    end
+
+    operators = OperatorEnum(;
+        default_params...,
+        binary_operators=(+, *, /, -, my_pretty_op),
+        unary_operators=(cos, sin),
+    )
+    @extend_operators operators
+
+    x1, x2 = [Node(; feature=i) for i in 1:2]
+
+    # Test default format (not pretty)
+    tree = my_pretty_op(x1, x2)
+    @test string_tree(tree, operators) == "my_pretty_op(x1, x2)"
+
+    # Test pretty format
+    @test string_tree(tree, operators; pretty=true) == "pretty_op_two(x1, x2)"
+
+    # Test with nested expressions
+    tree = sin(my_pretty_op(x1, x2))
+    @test string_tree(tree, operators) == "sin(my_pretty_op(x1, x2))"
+    @test string_tree(tree, operators; pretty=true) == "sin(pretty_op_two(x1, x2))"
+
+    # Test with constants
+    tree = my_pretty_op(x1, Node(; val=3.14))
+    @test string_tree(tree, operators) == "my_pretty_op(x1, 3.14)"
+    @test string_tree(tree, operators; pretty=true) == "pretty_op_two(x1, 3.14)"
+
+    # Test that the default implementation of get_pretty_op_name falls back to get_op_name
+    tree = sin(x1)
+    @test string_tree(tree, operators) == "sin(x1)"
+    @test string_tree(tree, operators; pretty=true) == "sin(x1)"
+
+    # Test with a unary operator that has a different pretty name
+    @eval begin
+        my_unary_op(x) = sin(x)
+        DE.get_op_name(::typeof(my_unary_op)) = "my_unary_op"
+        DE.get_pretty_op_name(::typeof(my_unary_op)) = "sine"
+    end
+
+    operators_with_unary = OperatorEnum(;
+        default_params...,
+        binary_operators=(+, *, /, -),
+        unary_operators=(cos, sin, my_unary_op),
+    )
+    @extend_operators operators_with_unary
+
+    tree = my_unary_op(x1)
+    @test string_tree(tree, operators_with_unary) == "my_unary_op(x1)"
+    @test string_tree(tree, operators_with_unary; pretty=true) == "sine(x1)"
+end