Merge pull request #192 from JuliaSymbolics/sm/print

Fix printing, similarterm, and add tree printing

Author: shashi

Project.toml

@@ -1,10 +1,11 @@
 name = "SymbolicUtils"
 uuid = "d1185830-fcd6-423d-90d6-eec64667417b"
 authors = ["Shashi Gowda"]
-version = "0.8.0"
+version = "0.8.1"
 
 [deps]
 AbstractAlgebra = "c3fe647b-3220-5bb0-a1ea-a7954cac585d"
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
@@ -17,6 +18,7 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
 AbstractAlgebra = "0.9, 0.10, 0.11, 0.12"
+AbstractTrees = "0.3"
 Combinatorics = "1.0"
 DataStructures = "0.18"
 IfElse = "0.1"

src/types.jl

@@ -370,7 +370,7 @@ function show_add(io, args)
             print_arg(io, -, t)
         else
             print(io, " - ")
-            print_arg(io, +, -t)
+            print_arg(IOContext(io, :paren=>true), +, -t)
         end
     end
 end
@@ -510,6 +510,11 @@ and the key (in Add) should instead be used to store the actual coefficient
 function makeadd(sign, coeff, xs...)
     d = sdict()
     for x in xs
+        if x isa Add
+            coeff += x.coeff
+            _merge!(+, d, x.dict, filter=_iszero)
+            continue
+        end
         if x isa Number
             coeff += x
             continue
@@ -631,7 +636,7 @@ function makemul(coeff, xs...; d=sdict())
             coeff *= x
         elseif x isa Mul
             coeff *= x.coeff
-            d = _merge(+, d, x.dict, filter=_iszero)
+            _merge!(+, d, x.dict, filter=_iszero)
         else
             v = 1 + get(d, x, 0)
             if _iszero(v)
@@ -721,8 +726,9 @@ end
 
 *(a::Pow, b::Mul) = b * a
 
-function _merge(f, d, others...; filter=x->false)
-    acc = copy(d)
+_merge(f, d, others...; filter=x->false) = _merge!(f, copy(d), others...; filter=filter)
+function _merge!(f, d, others...; filter=x->false)
+    acc = d
     for other in others
         for (k, v) in other
             v = f(v)
@@ -770,3 +776,27 @@ function Base.hash(t::Union{Add,Mul}, u::UInt64)
     t.hash[] = h′
     return h′
 end
+
+import AbstractTrees
+
+struct TreePrint
+    op
+    x
+end
+AbstractTrees.children(x::Term) = arguments(x)
+AbstractTrees.children(x::Union{Add, Mul}) = map(y->TreePrint(x isa Add ? (:*) : (:^), y), collect(pairs(x.dict)))
+AbstractTrees.children(x::Union{Pow}) = [x.base, x.exp]
+AbstractTrees.children(x::TreePrint) = [x.x[1], x.x[2]]
+
+print_tree(x; maxdepth=Inf, kw...) = print_tree(stdout, x; maxdepth=maxdepth, kw...)
+function print_tree(_io::IO, x::Union{Term, Add, Mul, Pow}; kw...)
+    AbstractTrees.print_tree(_io, x; withinds=true, kw...) do io, y, inds
+        if istree(y)
+            print(io, operation(y))
+        elseif y isa TreePrint
+            print(io, "(", y.op, ")")
+        else
+            print(io, y)
+        end
+    end
+end

test/basics.jl

@@ -111,6 +111,13 @@ end
     @test repr(Term(*, [1, 1])) == "*1"
     @test repr(Term(*, [2, 1])) == "2*1"
     @test repr((a + b) - (b + c)) == "a - c"
+    @test repr(a + -1*(b + c)) == "a - (b + c)"
+    @test repr(a + -1*b) == "a - b"
+end
+
+@testset "similarterm with Add" begin
+    @syms a b c
+    @test isequal(SymbolicUtils.similarterm((b + c), +, [a,  (b+c)]).dict, Dict(a=>1,b=>1,c=>1))
 end
 
 toterm(t) = Term{symtype(t)}(operation(t), arguments(t))