support output assignment for function calls

oqpy/program.py

@@ -423,11 +423,18 @@ def barrier(self, qubits_or_frames: Iterable[AstConvertible]) -> Program:
         self._add_statement(ast.QuantumBarrier(ast_qubits_or_frames))
         return self
 
-    def function_call(self, name: str, args: Iterable[AstConvertible]) -> None:
-        """Add a function call."""
-        self._add_statement(
-            ast.ExpressionStatement(ast.FunctionCall(ast.Identifier(name), map_to_ast(self, args)))
-        )
+    def function_call(
+        self,
+        name: str,
+        args: Iterable[AstConvertible],
+        assigns_to: AstConvertible = None,
+    ) -> None:
+        """Add a function call with an optional output assignment."""
+        function_call_node = ast.FunctionCall(ast.Identifier(name), map_to_ast(self, args))
+        if assigns_to is None:
+            self.do_expression(function_call_node)
+        else:
+            self._do_assignment(to_ast(self, assigns_to), "=", function_call_node)
 
     def play(self, frame: AstConvertible, waveform: AstConvertible) -> Program:
         """Play a waveform on a particular frame."""

tests/test_directives.py

@@ -2300,3 +2300,26 @@ def test_qubit_array():
         prog = oqpy.Program()
         q = oqpy.Qubit("q")
         prog.gate(q[0], "h")
+
+
+@pytest.mark.parametrize(
+    "args,assigns_to,expected",
+    [
+        ([], None, "OPENQASM 3.0;\nmy_function();"),
+        (
+            [oqpy.BitVar(name="a0"), oqpy.BitVar(name="a1")],
+            None,
+            "OPENQASM 3.0;\nbit a0;\nbit a1;\nmy_function(a0, a1);",
+        ),
+        (
+            [oqpy.BitVar(name="a0")],
+            oqpy.BitVar(name="b0"),
+            "OPENQASM 3.0;\nbit a0;\nbit b0;\nb0 = my_function(a0);",
+        ),
+    ],
+)
+def test_function_call(args, assigns_to, expected):
+    prog = Program()
+    prog.function_call("my_function", args, assigns_to)
+    assert prog.to_qasm() == expected
+    _check_respects_type_hints(prog)