Respond to CR

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:

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
 
 

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)

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()