# 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.
"""Operations to override return statement behavior."""
from collections.abc import Iterable
from typing import Any
from autoqasm import program
from autoqasm import types as aq_types
[docs]
def return_output_from_main(name: str, value: Any) -> Any:
"""Registers an output variable in the program with the given name that matches the type
of the value. The value is assigned to the output variable.
Args:
name (str): The symbol name for the return variable.
value (Any): The value of the returned variable.
Returns:
Any: Returns the same value that is being returned in the statement.
"""
aq_context = program.get_program_conversion_context()
if isinstance(value, Iterable):
for i, item in enumerate(value):
aq_context.register_output_parameter(f"{name}{i}", item)
return aq_types.wrap_value(tuple(value))
else:
aq_context.register_output_parameter(name, value)
return aq_types.wrap_value(value)