diff --git a/mapytex/calculus/core/tree.py b/mapytex/calculus/core/tree.py index 7a1f13c..f228b51 100644 --- a/mapytex/calculus/core/tree.py +++ b/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()} + {} + """ + 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