feat: allow partial updates to with_metadata

Author: MilesCranmer

src/Expression.jl

@@ -31,22 +31,25 @@ import ..SimplifyModule: combine_operators, simplify_tree!
 struct Metadata{NT<:NamedTuple}
     _data::NT
 end
-_data(x::Metadata) = getfield(x, :_data)
+unpack_metadata(x) = x  # Fallback for when the user doesn't use the Metadata type
+unpack_metadata(x::Metadata) = getfield(x, :_data)
 
-Base.propertynames(x::Metadata) = propertynames(_data(x))
-@unstable @inline Base.getproperty(x::Metadata, f::Symbol) = getproperty(_data(x), f)
-Base.show(io::IO, x::Metadata) = print(io, "Metadata(", _data(x), ")")
+Base.propertynames(x::Metadata) = propertynames(unpack_metadata(x))
+@unstable @inline function Base.getproperty(x::Metadata, f::Symbol)
+    return getproperty(unpack_metadata(x), f)
+end
+Base.show(io::IO, x::Metadata) = print(io, "Metadata(", unpack_metadata(x), ")")
 @inline _copy(x) = copy(x)
 @inline _copy(x::NamedTuple) = copy_named_tuple(x)
 @inline _copy(x::Nothing) = nothing
 @inline function copy_named_tuple(nt::NamedTuple)
     return NamedTuple{keys(nt)}(map(_copy, values(nt)))
 end
 @inline function Base.copy(metadata::Metadata)
-    return Metadata(_copy(_data(metadata)))
+    return Metadata(_copy(unpack_metadata(metadata)))
 end
-@inline Base.:(==)(x::Metadata, y::Metadata) = _data(x) == _data(y)
-@inline Base.hash(x::Metadata, h::UInt) = hash(_data(x), h)
+@inline Base.:(==)(x::Metadata, y::Metadata) = unpack_metadata(x) == unpack_metadata(y)
+@inline Base.hash(x::Metadata, h::UInt) = hash(unpack_metadata(x), h)
 
 """
     AbstractExpression{T,N}
@@ -216,7 +219,9 @@ end
 Create a new expression based on `ex` but with a different `metadata`.
 """
 function with_metadata(ex::AbstractExpression; metadata...)
-    return with_metadata(ex, Metadata((; metadata...)))
+    return with_metadata(
+        ex, Metadata((; unpack_metadata(get_metadata(ex))..., metadata...))
+    )
 end
 function with_metadata(ex::AbstractExpression, metadata::Metadata)
     return constructorof(typeof(ex))(get_contents(ex), metadata)

src/ParametricExpression.jl

@@ -5,7 +5,8 @@ using ChainRulesCore: ChainRulesCore as CRC, NoTangent, @thunk
 
 using ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum
 using ..NodeModule: AbstractExpressionNode, Node, tree_mapreduce
-using ..ExpressionModule: AbstractExpression, Metadata, with_contents, with_metadata
+using ..ExpressionModule:
+    AbstractExpression, Metadata, with_contents, with_metadata, unpack_metadata
 using ..ChainRulesModule: NodeTangent
 
 import ..NodeModule:
@@ -63,7 +64,6 @@ mutable struct ParametricNode{T} <: AbstractExpressionNode{T}
         return n
     end
 end
-@inline _data(x::Metadata) = getfield(x, :_data)
 
 """
     ParametricExpression{T,N<:ParametricNode{T},D<:NamedTuple} <: AbstractExpression{T,N}
@@ -79,7 +79,9 @@ struct ParametricExpression{
     metadata::Metadata{D}
 
     function ParametricExpression(tree::ParametricNode, metadata::Metadata)
-        return new{eltype(tree),typeof(tree),typeof(_data(metadata))}(tree, metadata)
+        return new{eltype(tree),typeof(tree),typeof(unpack_metadata(metadata))}(
+            tree, metadata
+        )
     end
 end
 function ParametricExpression(

src/StructuredExpression.jl

@@ -16,7 +16,7 @@ import ..ExpressionModule:
     with_contents,
     Metadata,
     _copy,
-    _data,
+    unpack_metadata,
     default_node_type,
     node_type,
     get_scalar_constants,
@@ -114,7 +114,7 @@ constructorof(::Type{<:StructuredExpression}) = StructuredExpression
 function Base.copy(e::AbstractStructuredExpression)
     ts = get_contents(e)
     meta = get_metadata(e)
-    meta_inner = _data(meta)
+    meta_inner = unpack_metadata(meta)
     copy_ts = NamedTuple{keys(ts)}(map(copy, values(ts)))
     keys_except_structure = filter(!=(:structure), keys(meta_inner))
     copy_metadata = (;

test/test_expressions.jl

@@ -413,3 +413,26 @@ end
 
     #literate_end
 end
+
+@testitem "Expression with_metadata partial updates" begin
+    using DynamicExpressions
+    using DynamicExpressions: get_operators, get_metadata, with_metadata, get_variable_names
+
+    # Create an expression with initial metadata
+    ex = @parse_expression(
+        x1 + 1.5,
+        operators = OperatorEnum(; binary_operators=[+, *]),
+        variable_names = ["x1"]
+    )
+
+    # Update only the variable_names, keeping the original operators
+    new_ex = with_metadata(ex; variable_names=["y1"])
+    @test get_variable_names(new_ex, nothing) == ["y1"]
+    @test get_operators(new_ex, nothing) == get_operators(ex, nothing)
+
+    # Update only the operators, keeping the original variable_names
+    new_operators = OperatorEnum(; binary_operators=[+])
+    new_ex2 = with_metadata(ex; operators=new_operators)
+    @test get_variable_names(new_ex2, nothing) == ["x1"]
+    @test get_operators(new_ex2, nothing) == new_operators
+end