Source code for autoqasm.converters.return_statements

# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.


"""Converters for return statement nodes."""

import ast

from malt.converters import return_statements
from malt.core import ag_ctx, converter
from malt.pyct import templates

from autoqasm import constants, program
from autoqasm.operators.assignments import assign_for_output




class ReturnTransformer(converter.Base):


    def visit_Return(self, node: ast.stmt) -> ast.stmt:
        """AutoQASM-specific return statement transformations.

        Args:
            node (ast.stmt): Return statement node to transform.

        Returns:
            ast.stmt: Transformed return statement node.
        """
        aq_context = program.get_program_conversion_context()
        if aq_context.subroutines_processing or node.value is None:
            return node

        template = (
            f"name_ = ag__.{assign_for_output.__name__}(name_const_, "
            "ag__.return_output_from_main(name_const_, value_))"
        )

        name = constants.MAIN_RETURN_VAL_NAME
        if isinstance(node.value, ast.Name):
            name = node.value.id

        return templates.replace(
            template,
            name_=name,
            name_const_=ast.Constant(name, None),
            value_=node.value,
            original=node,
        )






def transform(
    node: ast.stmt, ctx: ag_ctx.ControlStatusCtx, default_to_null_return: bool = True
) -> ast.stmt:
    """Handle AutoQASM-specific return statement functionality before
    passing control to AutoGraph.

    Args:
        node (ast.stmt): AST node to transform.
        ctx (ag_ctx.ControlStatusCtx): Transformer context.
        default_to_null_return (bool): Configuration option.

    Returns:
        ast.stmt: Transformed node.
    """
    node = ReturnTransformer(ctx).visit(node)
    return return_statements.transform(node, ctx)