Allow latest variable names set

Author: MilesCranmer

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
-version = "0.12.0"
+version = "0.12.1"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/OperatorEnumConstruction.jl

@@ -27,17 +27,22 @@ const ALREADY_DEFINED_UNARY_OPERATORS = (;
 const ALREADY_DEFINED_BINARY_OPERATORS = (;
     operator_enum=Dict{Function,Bool}(), generic_operator_enum=Dict{Function,Bool}()
 )
+const LATEST_VARIABLE_NAMES = Ref{Union{Nothing,Vector{String}}}(nothing)
 
 function Base.show(io::IO, tree::Node)
     latest_operators_type = LATEST_OPERATORS_TYPE.x
     if latest_operators_type == IsNothing
-        return print(io, string_tree(tree))
+        return print(io, string_tree(tree; variable_names=LATEST_VARIABLE_NAMES.x))
     elseif latest_operators_type == IsOperatorEnum
         latest_operators = LATEST_OPERATORS.x::OperatorEnum
-        return print(io, string_tree(tree, latest_operators))
+        return print(
+            io, string_tree(tree, latest_operators; variable_names=LATEST_VARIABLE_NAMES.x)
+        )
     else
         latest_operators = LATEST_OPERATORS.x::GenericOperatorEnum
-        return print(io, string_tree(tree, latest_operators))
+        return print(
+            io, string_tree(tree, latest_operators; variable_names=LATEST_VARIABLE_NAMES.x)
+        )
     end
 end
 function (tree::Node)(X; kws...)

test/test_print.jl

@@ -98,3 +98,16 @@ end
     f_constant(val::Float64, args...) = string(round(val; digits=4))
     @test string_tree(tree, operators; f_constant=f_constant) == "((x1 * x1) + 3.1416)"
 end
+
+@testset "Test variable names" begin
+    operators = OperatorEnum(;
+        binary_operators=[+, *, /, -], unary_operators=[cos, sin],
+    )
+    @extend_operators operators
+    x1, x2, x3 = [Node(Float64; feature=i) for i in 1:3]
+    DynamicExpressions.OperatorEnumConstructionModule.LATEST_VARIABLE_NAMES.x = ["k1", "k2", "k3"]
+    tree = x1 * x2 + x3
+    @test string(tree) == "((k1 * k2) + k3)"
+    DynamicExpressions.OperatorEnumConstructionModule.LATEST_VARIABLE_NAMES.x = nothing
+    @test string(tree) == "((x1 * x2) + x3)"
+end