Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# coding: utf-8
2"""
3Helpers for :epkg:`sympy`.
4"""
7def enumerate_traverse(e, parent=None):
8 """
9 Enumerates all nodes in *e*.
11 :param e: :epkg:`sympy` expression or node
12 :param parent: parent of *e*
13 :return: iterator on itself and the children
14 """
15 yield dict(e=e, p=parent)
16 for arg in e.args:
17 for r in enumerate_traverse(arg, e):
18 yield r
21class SympyNode:
22 """
23 Model a :epkg:`sympy` expression.
24 It probably exist in :epkg:`sympy`.
26 :param e: :epkg:`sympy` expression
27 :param parent: *SympyNode*
28 """
30 def __init__(self, e, parent=None):
31 if parent is not None and not isinstance(parent, SympyNode):
32 raise TypeError( # pragma: no cover
33 "parent must be None or SympyNode")
34 self._e = e
35 self._parent = parent
36 self._children = []
37 for a in self._e.args:
38 self._children.append(SympyNode(a, self))
40 def __repr__(self):
41 return 'SympyNode("%r")' % self.element
43 @property
44 def element(self):
45 return self._e
47 @property
48 def parent(self):
49 return self._parent
51 @property
52 def children(self):
53 return self._children
55 def __iter__(self):
56 """
57 Enumerate all nodes.
58 """
59 yield self.element
60 for a in self.children:
61 for n in a:
62 yield a
64 def enumerate_parents(self):
65 """
66 Enumerates all parents including itself.
67 """
68 n = self
69 while n.parent is not None:
70 yield n
71 n = n.parent