| 86 | |
| 87 | |
| 88 | class Node(object): |
| 89 | def __init__(self, child_list=None, data=None, name="unnamed", arg_name=""): |
| 90 | """ |
| 91 | Hold a function call with its bound arguments, along with child nodes. |
| 92 | |
| 93 | """ |
| 94 | self.parent = None |
| 95 | self.name = name |
| 96 | self.arg_name = arg_name |
| 97 | self.child_list = [] if child_list is None else child_list |
| 98 | self.data = {} if data is None else data |
| 99 | self.updated = False |
| 100 | # hacky way to add their argument name when a function of tests was given |
| 101 | |
| 102 | def __call__(self, state=None): |
| 103 | """Call original function with its arguments, and optional state""" |
| 104 | ba = self.data["bound_args"] |
| 105 | if state: |
| 106 | func = self.data["func"] |
| 107 | ba = inspect.signature(func).bind(state, *ba.args[1:], **ba.kwargs) |
| 108 | func(*ba.args, **ba.kwargs) |
| 109 | return state |
| 110 | else: |
| 111 | self.data["func"](*ba.args, **ba.kwargs) |
| 112 | ba.apply_defaults() |
| 113 | return ba.arguments["state"] |
| 114 | |
| 115 | def __str__(self): |
| 116 | return pp.pformat( |
| 117 | dict(getattr(self.data.get("bound_args", {}), "arguments", {})) |
| 118 | ) |
| 119 | |
| 120 | def __iter__(self): |
| 121 | for c in self.child_list: |
| 122 | yield c |
| 123 | |
| 124 | def partial(self): |
| 125 | """Return partial of original function call""" |
| 126 | ba = self.data["bound_args"] |
| 127 | return state_partial(self.data["func"], *ba.args[1:], **ba.kwargs) |
| 128 | |
| 129 | def update_child_calls(self): |
| 130 | """Replace child nodes on original function call with their partials""" |
| 131 | |
| 132 | for node in filter(lambda n: len(n.arg_name), self.child_list): |
| 133 | self.data["bound_args"].arguments[node.arg_name] = node.partial() |
| 134 | self.updated = True |
| 135 | |
| 136 | def remove_child(self, node): |
| 137 | index = self.child_list.index(node) |
| 138 | del self.child_list[index] |
| 139 | return index |
| 140 | |
| 141 | def add_child(self, child): |
| 142 | # since it is a tree, there is only one parent |
| 143 | # note this means we do not allow edges between same layer units |
| 144 | if child.parent: |
| 145 | child.parent.remove_child(child) |