traverse.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637
  1. """Strategies to Traverse a Tree."""
  2. from sympy.strategies.util import basic_fns
  3. from sympy.strategies.core import chain, do_one
  4. def top_down(rule, fns=basic_fns):
  5. """Apply a rule down a tree running it on the top nodes first."""
  6. return chain(rule, lambda expr: sall(top_down(rule, fns), fns)(expr))
  7. def bottom_up(rule, fns=basic_fns):
  8. """Apply a rule down a tree running it on the bottom nodes first."""
  9. return chain(lambda expr: sall(bottom_up(rule, fns), fns)(expr), rule)
  10. def top_down_once(rule, fns=basic_fns):
  11. """Apply a rule down a tree - stop on success."""
  12. return do_one(rule, lambda expr: sall(top_down(rule, fns), fns)(expr))
  13. def bottom_up_once(rule, fns=basic_fns):
  14. """Apply a rule up a tree - stop on success."""
  15. return do_one(lambda expr: sall(bottom_up(rule, fns), fns)(expr), rule)
  16. def sall(rule, fns=basic_fns):
  17. """Strategic all - apply rule to args."""
  18. op, new, children, leaf = map(fns.get, ('op', 'new', 'children', 'leaf'))
  19. def all_rl(expr):
  20. if leaf(expr):
  21. return expr
  22. else:
  23. args = map(rule, children(expr))
  24. return new(op(expr), *args)
  25. return all_rl