Repository: onnx/onnx
File: onnx/utils.py (Python)

Code before:
import os
import tarfile

import onnx.checker
import onnx.helper
import onnx.shape_inference
from onnx import FunctionProto, ModelProto, NodeProto, TensorProto, ValueInfoProto


class Extractor:
    def __init__(self, model: ModelProto) -> None:
        self.model = model
        self.graph = self.model.graph
        self.wmap = self._build_name2obj_dict(self.graph.initializer)
        self.vimap = self._build_name2obj_dict(self.graph.value_info)

    @staticmethod
    def _build_name2obj_dict(objs):  # type: ignore
        return {obj.name: obj for obj in objs}

    def _collect_new_io_core(
        self, original_io, io_names_to_extract
    ) -> list[ValueInfoProto]:
        original_io_map = self._build_name2obj_dict(original_io)
        new_io_tensors = []
        for io_name_to_extract in io_names_to_extract:  # assume no duplicates
            if io_name_to_extract in original_io_map:
                new_io_tensors.append(original_io_map[io_name_to_extract])
            else:
                new_io_tensors.append(self.vimap[io_name_to_extract])
        return new_io_tensors  # same order as io_names_to_extract

    def _collect_new_inputs(self, names: list[str]) -> list[ValueInfoProto]:
        return self._collect_new_io_core(self.graph.input, names)

    def _collect_new_outputs(self, names: list[str]) -> list[ValueInfoProto]:
        return self._collect_new_io_core(self.graph.output, names)

    def _dfs_search_reachable_nodes(
        self,
        node_output_name: str,
        graph_input_names: set[str],
        nodes: list[NodeProto],
        reachable: set[int],
        unreachable: set[int],
    ) -> None:
        """Helper function to find nodes which are connected to an output

        Arguments:
            node_output_name (str): The name of the output
            graph_input_names (set of string): The names of all inputs of the graph

Review comment (refactor):
```suggestion
        for io_name_to_extract in io_names_to_extract:
```

Code after:
import os
import tarfile

from google.protobuf.internal.containers import RepeatedCompositeFieldContainer

import onnx.checker
import onnx.helper
import onnx.shape_inference
from onnx import FunctionProto, ModelProto, NodeProto, TensorProto, ValueInfoProto


class Extractor:
    def __init__(self, model: ModelProto) -> None:
        self.model = model
        self.graph = self.model.graph
        self.wmap = self._build_name2obj_dict(self.graph.initializer)
        self.vimap = self._build_name2obj_dict(self.graph.value_info)

    @staticmethod
    def _build_name2obj_dict(objs):  # type: ignore
        return {obj.name: obj for obj in objs}

    def _collect_new_io_core(
        self,
        original_io: RepeatedCompositeFieldContainer[ValueInfoProto],
        io_names_to_extract: list[str],
    ) -> list[ValueInfoProto]:
        original_io_map = self._build_name2obj_dict(original_io)
        new_io_tensors = []
        for io_name_to_extract in io_names_to_extract:
            if io_name_to_extract in original_io_map:
                new_io_tensors.append(original_io_map[io_name_to_extract])
            else:
                new_io_tensors.append(self.vimap[io_name_to_extract])
        return new_io_tensors  # same order as io_names_to_extract

    def _collect_new_inputs(self, names: list[str]) -> list[ValueInfoProto]:
        return self._collect_new_io_core(self.graph.input, names)

    def _collect_new_outputs(self, names: list[str]) -> list[ValueInfoProto]:
        return self._collect_new_io_core(self.graph.output, names)

    def _dfs_search_reachable_nodes(
        self,
        node_output_name: str,
        graph_input_names: set[str],
        nodes: list[NodeProto],
        reachable: set[int],
        unreachable: set[int],
    ) -> None:
        """Helper function to find nodes which are connected to an output