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

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