Coverage for aftercovid/models/_sympy_helper.py: 100%
35 statements
« prev ^ index » next coverage.py v7.1.0, created at 2024-03-29 03:09 +0100
« prev ^ index » next coverage.py v7.1.0, created at 2024-03-29 03:09 +0100
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 f'SympyNode("{self.element!r}")'
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