def generate_code()

in core/maxframe/codegen.py [0:0]

• 38 lines of code
• 13 McCabe index (conditional complexity)

    def generate_code(self, dag: TileableGraph) -> List[str]:
        """
        Generate the code of the input dag.

        Parameters
        ----------
        dag : TileableGraph
            The input DAG instance.

        Returns
        -------
        List[str] :
            The code lines.
        """
        code_lines = []
        visited_op_key = set()
        result_key_set = set(t.key for t in dag.result_tileables)
        out_refcounts = dict()
        for tileable in dag.topological_iter():
            op: OperatorType = tileable.op
            if op.key in visited_op_key or isinstance(op, Fetch):
                continue

            visited_op_key.add(op.key)

            adapter = self.get_op_adapter(type(op))()
            code_lines.extend(adapter.generate_pre_op_code(op, self._context))
            if self._generate_comments_enabled:
                code_lines.extend(adapter.generate_comment(op, self._context))
            code_lines.extend(adapter.generate_code(op, self._context))
            code_lines.extend(adapter.generate_post_op_code(op, self._context))
            code_lines.append("")  # Append an empty line to separate operators

            # record refcounts
            for out_t in op.outputs:
                if out_t.key in result_key_set:
                    continue
                if dag.count_successors(out_t) == 0:
                    delete_code = self._generate_delete_code(
                        self._context.get_input_tileable_variable(out_t)
                    )
                    code_lines.extend(delete_code)
                else:
                    out_refcounts[out_t.key] = dag.count_successors(out_t)

            # check if refs of inputs are no longer needed
            for inp_t in op.inputs:
                if inp_t.key not in out_refcounts:
                    continue
                out_refcounts[inp_t.key] -= 1
                if out_refcounts[inp_t.key] == 0:
                    delete_code = self._generate_delete_code(
                        self._context.get_input_tileable_variable(inp_t)
                    )
                    code_lines.extend(delete_code)
                    out_refcounts.pop(inp_t.key)

        return code_lines