| 118 | |
| 119 | |
| 120 | class ArgsTree(object): |
| 121 | def __init__( |
| 122 | self, |
| 123 | io_args: Union[Tuple, List, Dict], |
| 124 | gen_name: bool = False, |
| 125 | root_prefix: str = "", |
| 126 | root_name: str = None, |
| 127 | tensor_type=Tensor, |
| 128 | ) -> None: |
| 129 | |
| 130 | self._io_args = io_args |
| 131 | self._gen_name = gen_name |
| 132 | self._root_prefix = root_prefix |
| 133 | self._root_name = root_name |
| 134 | self._named_io_args = None |
| 135 | self._next_global_index = 0 |
| 136 | self._tensor_type = tensor_type |
| 137 | |
| 138 | if self._gen_name: |
| 139 | self._named_io_args = self._construct_named_io_args( |
| 140 | self._io_args, self._root_prefix, self._root_name |
| 141 | ) |
| 142 | |
| 143 | def gen_name(self): |
| 144 | return self._gen_name |
| 145 | |
| 146 | def iter_nodes(self): |
| 147 | r""" |
| 148 | return a generator of the args tree nodes in the DFS manner. |
| 149 | The node returned can be of type NamedArg or non-NamedArg depending on whether gen_name is set. |
| 150 | If gen_name is set, the node will be NamedArg. |
| 151 | """ |
| 152 | |
| 153 | if self._gen_name: |
| 154 | args_to_iter = self._named_io_args |
| 155 | else: |
| 156 | args_to_iter = self._io_args |
| 157 | |
| 158 | # NOTE(lixiang): Generator expression and iterator are used. |
| 159 | # This avoids generating the full list in memory and only processes the nodes that need to be processed, |
| 160 | # reducing time and space consumption. |
| 161 | stack = [iter([args_to_iter])] |
| 162 | while len(stack) > 0: |
| 163 | try: |
| 164 | curr = next(stack[-1]) |
| 165 | if _is_raw_type(curr, NamedArg): |
| 166 | curr_value = curr.value() |
| 167 | else: |
| 168 | curr_value = curr |
| 169 | |
| 170 | if _is_raw_type(curr_value, list) or _is_raw_type(curr_value, tuple): |
| 171 | children = curr_value |
| 172 | elif _is_raw_type(curr_value, dict) or _is_raw_type( |
| 173 | curr_value, OrderedDict |
| 174 | ): |
| 175 | children = curr_value.values() |
| 176 | else: |
| 177 | children = None |
no outgoing calls