(key, steps_or_nets, target)
| 100 | |
| 101 | |
| 102 | def get_setup_nets(key, steps_or_nets, target): |
| 103 | init_net = core.Net(key + '/init') |
| 104 | exit_net = core.Net(key + '/exit') |
| 105 | init_nets = [] |
| 106 | exit_nets = [] |
| 107 | objs = [] |
| 108 | for step_or_net in steps_or_nets: |
| 109 | if hasattr(step_or_net, 'get_all_attributes'): |
| 110 | objs += step_or_net.get_all_attributes(key) |
| 111 | elif hasattr(step_or_net, 'get_attributes'): |
| 112 | objs += step_or_net.get_attributes(key) |
| 113 | for obj in objs: |
| 114 | # these are needed in order to allow nesting of TaskGroup, which |
| 115 | # is a feature not yet implemented. |
| 116 | if hasattr(obj, '_setup_used') and obj._setup_used: |
| 117 | continue |
| 118 | if hasattr(obj, '_setup_target') and obj._setup_target != target: |
| 119 | continue |
| 120 | if hasattr(obj, 'setup'): |
| 121 | nets = obj.setup(init_net) |
| 122 | if isinstance(nets, (list, tuple)): |
| 123 | init_nets += nets |
| 124 | elif isinstance(nets, (core.Net, core.ExecutionStep)): |
| 125 | init_nets.append(nets) |
| 126 | elif nets is not None: |
| 127 | raise TypeError('Unsupported type for setup: %s' % type(nets)) |
| 128 | obj._setup_used = True |
| 129 | if hasattr(obj, 'exit'): |
| 130 | nets = obj.exit(exit_net) |
| 131 | if isinstance(nets, (list, tuple)): |
| 132 | exit_nets += nets |
| 133 | elif isinstance(nets, (core.Net, core.ExecutionStep)): |
| 134 | exit_nets.append(nets) |
| 135 | elif nets is not None: |
| 136 | raise TypeError('Unsupported type for setup: %s' % type(nets)) |
| 137 | obj._setup_used = True |
| 138 | |
| 139 | if len(init_net.Proto().op) > 0: |
| 140 | init_nets.insert(0, init_net) |
| 141 | if len(exit_net.Proto().op) > 0: |
| 142 | exit_nets.insert(0, exit_net) |
| 143 | return init_nets, exit_nets |
| 144 | |
| 145 | |
| 146 | def add_setup_steps(step, init_nets, exit_nets, name): |
no test coverage detected
searching dependent graphs…