get leafs or operator methodes for tree

mapytex/calculus/core/tree.py

@@ -353,6 +353,64 @@ class Tree(object):
         """
         return self.apply(postfix_concatenate)
 
+    def get_leafs(self, callback = lambda x:x):
+        """ Generator which yield all the leaf value of the tree.
+        Callback act on every leaf.
+
+        :param callback: function on leaf
+
+        :example:
+
+        >>> nested_par = ("+", (
+        ...    ("*", (3, 4)),
+        ...    2))
+        >>> t = Tree.from_nested_parenthesis(nested_par)
+        >>> [l for l in t.get_leafs()]
+        [3, 4, 2]
+        >>> {type(l) for l in t.get_leafs()}
+        {<class 'int'>}
+        """
+        try:
+            yield from self.left_value.get_leafs(callback)
+        except AttributeError:
+            yield callback(self.left_value)
+        try:
+            yield from self.right_value.get_leafs(callback)
+        except AttributeError:
+            yield callback(self.right_value)
+
+    def get_nodes(self, callback = lambda x:x):
+        """ Generator which yield all nodes of the tree.
+        Callback act on every nodes.
+
+        :param callback: function on node
+
+        :example:
+
+        >>> nested_par = ("+", (
+        ...    ("*", (3, 4)),
+        ...    2))
+        >>> t = Tree.from_nested_parenthesis(nested_par)
+        >>> [l for l in t.get_nodes()]
+        ['+', '*']
+        >>> nested_par = ("+", (
+        ...    ("*", (3, 4)),
+        ...    ("*", (3, 4))
+        ...    ))
+        >>> t = Tree.from_nested_parenthesis(nested_par)
+        >>> [l for l in t.get_nodes()]
+        ['+', '*', '*']
+        """
+        yield self.node
+        try:
+            yield from self.left_value.get_nodes(callback)
+        except AttributeError:
+            pass
+        try:
+            yield from self.right_value.get_nodes(callback)
+        except AttributeError:
+            pass
+
     def __str__(self):
         """ Overload str method