diff --git a/src/braket/experimental/autoqasm/api.py b/src/braket/experimental/autoqasm/api.py index 5c912a48c..722c6b7d8 100644 --- a/src/braket/experimental/autoqasm/api.py +++ b/src/braket/experimental/autoqasm/api.py @@ -245,16 +245,12 @@ def _convert_program_as_main( args (List[Any]): Arguments passed to the program when called. kwargs (Dict[str, Any]): Keyword arguments passed to the program when called. """ - program_conversion_context.in_main(True) - # Process the program aq_transpiler.converted_call(f, args, kwargs, options=options) # Modify program to add qubit declaration if necessary _add_qubit_declaration(program_conversion_context) - program_conversion_context.in_main(False) - def _convert_program_as_subroutine( f: Callable, @@ -276,8 +272,6 @@ def _convert_program_as_subroutine( args (List[Any]): Arguments passed to the program when called. kwargs (Dict[str, Any]): Keyword arguments passed to the program when called. """ - program_conversion_context.in_main(False) - oqpy_program = program_conversion_context.get_oqpy_program() if f not in program_conversion_context.subroutines_processing: diff --git a/src/braket/experimental/autoqasm/converters/return_statements.py b/src/braket/experimental/autoqasm/converters/return_statements.py index 2effc49b1..5089b4e98 100644 --- a/src/braket/experimental/autoqasm/converters/return_statements.py +++ b/src/braket/experimental/autoqasm/converters/return_statements.py @@ -25,8 +25,8 @@ class ReturnValidator(converter.Base): def visit_Return(self, node: ast.stmt) -> ast.stmt: aq_context = program.get_program_conversion_context() - if aq_context.is_in_main and node.value is not None: - warnings.warn("`output` is currently unsupported; `return` statement has no effect.") + if not aq_context.subroutines_processing and node.value is not None: + warnings.warn("Return value from top level function is ignored.") return node diff --git a/src/braket/experimental/autoqasm/program/program.py b/src/braket/experimental/autoqasm/program/program.py index f13edb034..9ededf4d4 100644 --- a/src/braket/experimental/autoqasm/program/program.py +++ b/src/braket/experimental/autoqasm/program/program.py @@ -88,7 +88,6 @@ def __init__(self, user_config: Optional[UserConfig] = None): self.return_variable = None self._qubits_seen = set() self._var_idx = 0 - self._in_main = False def make_program(self) -> Program: """Makes a Program object using the oqpy program from this conversion context. @@ -108,15 +107,6 @@ def qubits(self) -> List[int]: # Can be memoized or otherwise made more performant return sorted(list(self._qubits_seen)) - @property - def is_in_main(self) -> bool: - """Whether the main block is currently being processed.""" - return self._in_main - - def in_main(self, in_main: bool) -> None: - """Set context for processing the `main` block.""" - self._in_main = in_main - def register_qubit(self, qubit: int) -> None: """Register a virtual qubit to use in this program.""" self._qubits_seen.add(qubit) diff --git a/test/unit_tests/braket/experimental/autoqasm/test_api.py b/test/unit_tests/braket/experimental/autoqasm/test_api.py index 1ffa9564e..72def1948 100644 --- a/test/unit_tests/braket/experimental/autoqasm/test_api.py +++ b/test/unit_tests/braket/experimental/autoqasm/test_api.py @@ -846,7 +846,7 @@ def test_main_return(): def main() -> int: return 1 - with pytest.warns(UserWarning, match="`output` is currently unsupported"): + with pytest.warns(UserWarning, match="Return value from top level function is ignored"): main()