12345678910111213141516171819202122232425 |
- import functools
- import itertools
- # A generic depth-first tree-visit function that can let multiple visitor functions do "their thing" in only a single pass.
- # It accepts 2 lists of visitor functions:
- # 'parent_first' is a list of callbacks that will be called with 1 or 2 parameters (and therefore the 2nd parameter of the callback should have a default value): 1) the current element and 2) the value of the callback returned by the parent element if the current element is not the root element.
- # 'child_first' is a list of callbacks that will be called with 2 parameters: 2) the current element and 2) An empty list for all the leaf elements or a list with the responses of the children of that element for the callback.
- def visit_tree(node, get_children, parent_first=[], child_first=[], parent_values=None):
- # Most parameters unchanged for recursive calls
- def visit(node, parent_values):
- if parent_values is None:
- parent_values = [f(node) for f in parent_first]
- else:
- parent_values = [f(node, p) for f,p in zip(parent_first, parent_values)]
- # child_responses is a list of len(children)-many lists C, where for every child of node, C is the 'to_parent' value returned by every child
- child_responses = [visit(node=c, parent_values=parent_values) for c in get_children(node)]
- to_parent = [f(node, [cs[i] for cs in child_responses]) for i,f in enumerate(child_first)]
- # 'to_parent' is the mapping from our child_first-functions to those functions called on the responses we got from our children
- return to_parent
- return visit(node, parent_values)
|