Hide keyboard shortcuts

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""" 

5 

6 

7def enumerate_traverse(e, parent=None): 

8 """ 

9 Enumerates all nodes in *e*. 

10 

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 

19 

20 

21class SympyNode: 

22 """ 

23 Model a :epkg:`sympy` expression. 

24 It probably exist in :epkg:`sympy`. 

25 

26 :param e: :epkg:`sympy` expression 

27 :param parent: *SympyNode* 

28 """ 

29 

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

39 

40 def __repr__(self): 

41 return 'SympyNode("%r")' % self.element 

42 

43 @property 

44 def element(self): 

45 return self._e 

46 

47 @property 

48 def parent(self): 

49 return self._parent 

50 

51 @property 

52 def children(self): 

53 return self._children 

54 

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 

63 

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