Method iter_nodes

python/oneflow/framework/args_tree.py:146–185

r""" return a generator of the args tree nodes in the DFS manner. The node returned can be of type NamedArg or non-NamedArg depending on whether gen_name is set. If gen_name is set, the node will be NamedArg.

(self)

        return self._gen_name

    def iter_nodes(self):
        r"""
        return a generator of the args tree nodes in the DFS manner. 
        The node returned can be of type NamedArg or non-NamedArg depending on whether gen_name is set. 
        If gen_name is set, the node will be NamedArg. 
        """

        if self._gen_name:
            args_to_iter = self._named_io_args
        else:
            args_to_iter = self._io_args

        # NOTE(lixiang): Generator expression and iterator are used.
        #   This avoids generating the full list in memory and only processes the nodes that need to be processed,
        #   reducing time and space consumption.
        stack = [iter([args_to_iter])]
        while len(stack) > 0:
            try:
                curr = next(stack[-1])
                if _is_raw_type(curr, NamedArg):
                    curr_value = curr.value()
                else:
                    curr_value = curr

                if _is_raw_type(curr_value, list) or _is_raw_type(curr_value, tuple):
                    children = curr_value
                elif _is_raw_type(curr_value, dict) or _is_raw_type(
                    curr_value, OrderedDict
                ):
                    children = curr_value.values()
                else:
                    children = None

                if children:
                    stack.append(iter(children))

                yield curr

            except StopIteration:
                stack.pop()

    def iter_named_nodes(self):
        assert self._gen_name, "Only use this if gen_name is set!"