sympyexpression-treespostorder

sympy: post order traversal over expression tree: skips top-level operation


I am trying to loop post order through a sympy-expression. It works fine except for the fact that it skips the last addition (im my example case the addition of sympy.sin(x * y) and z**2).

import sympy

def post(expr):
    if not expr:
        return
    for arg in expr.args:
        post(arg)
        print(f'arg {arg}')
        print(f'arg.func: {arg.func}')
        print(f'arg.args: {arg.args}')

x, y, z = sympy.symbols('x,y,z')
expr = sympy.sin(x * y) + z**2
post(expr)

I think it's because of the for arg in expr.args part yet if I ditch the .args and loop over the whole expression

import sympy

def post(expr):
    if not expr:
        return
    for arg in expr:
        post(arg)
        print(f'arg {arg}')
        print(f'arg.func: {arg.func}')
        print(f'arg.args: {arg.args}')

x, y, z = sympy.symbols('x,y,z')
expr = sympy.sin(x * y) + z**2
post(expr)

I get TypeError: 'Add' object is not iterable. Can I somehow fix this such that it loops over all functions of the expression and doesn't skip the top-level function?


Solution

  • This battery is already included if you want to use it:

    from sympy.core.traversal import postorder_traversal
    for i in postorder_traversal(expr):print(i)
    

    gives

    z
    2
    z**2
    x
    y
    x*y
    sin(x*y)
    z**2 + sin(x*y)
    

    Otherwise, consider this modification for your function:

    def post(expr):
        for arg in expr.args:
            post(arg)
        print(f'arg {expr}')
        print(f'arg.func: {expr.func}')
        print(f'arg.args: {expr.args}')