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

Question

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?

Answer 1

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}')