Keep the arg names in OQFunctionCall (#75)

* complete test coverage for base.py

* simplify tests

* store the args dict in OQFunctionCall

* keep Iterable

* remove useless list()

* update type hint

Author: jcjaskula-aws

oqpy/classical_types.py

@@ -424,7 +424,7 @@ class OQFunctionCall(OQPyExpression):
     def __init__(
         self,
         identifier: Union[str, ast.Identifier],
-        args: Iterable[AstConvertible],
+        args: Union[Iterable[AstConvertible], dict[Any, AstConvertible]],
         return_type: Optional[ast.ClassicalType],
         extern_decl: ast.ExternDeclaration | None = None,
         subroutine_decl: ast.SubroutineDefinition | None = None,
@@ -433,7 +433,8 @@ def __init__(
 
         Args:
             identifier: The function name.
-            args: The function arguments.
+            args: The function arguments. If passed as a dict, the values are used when
+                creating the FunctionCall ast node.
             return_type: The type returned by the function call. If none, returns nothing.
             extern_decl: An optional extern declaration ast node. If present,
                 this extern declaration will be added to the top of the program
@@ -457,4 +458,5 @@ def to_ast(self, program: Program) -> ast.Expression:
             program.externs[self.identifier.name] = self.extern_decl
         if self.subroutine_decl is not None:
             program._add_subroutine(self.identifier.name, self.subroutine_decl)
-        return ast.FunctionCall(self.identifier, map_to_ast(program, self.args))
+        args = self.args.values() if isinstance(self.args, dict) else self.args
+        return ast.FunctionCall(self.identifier, map_to_ast(program, args))

oqpy/pulse.py

@@ -75,7 +75,9 @@ def __init__(
             init_expression = None
         else:
             assert frequency is not None
-            init_expression = OQFunctionCall("newframe", [port, frequency, phase], ast.FrameType)
+            init_expression = OQFunctionCall(
+                "newframe", {"port": port, "frequency": frequency, "phase": phase}, ast.FrameType
+            )
         super().__init__(
             init_expression, name, needs_declaration=needs_declaration, annotations=annotations
         )

oqpy/subroutines.py

@@ -153,7 +153,7 @@ def wrapper(
         program.externs.update(inner_prog.externs)
         return OQFunctionCall(
             identifier,
-            args,
+            {k: v for k, v in zip(argnames[1:], args)},
             return_type,
             subroutine_decl=stmt,
         )
@@ -218,7 +218,9 @@ def call_extern(*call_args: AstConvertible, **call_kwargs: AstConvertible) -> OQ
         for i, a in enumerate(call_args):
             if type(arg_types[i]) == ast.DurationType:
                 new_args[i] = convert_float_to_duration(a)
-        return OQFunctionCall(name, new_args, return_type, extern_decl=extern_decl)
+        return OQFunctionCall(
+            name, {k: v for k, v in zip(arg_names, new_args)}, return_type, extern_decl=extern_decl
+        )
 
     return call_extern
 

tests/test_directives.py

@@ -328,6 +328,7 @@ def test_non_trivial_array_access():
 
     assert prog.to_qasm() == expected
     _check_respects_type_hints(prog)
+    assert expr_matches(frame.init_expression.args["port"], port)
 
 
 def test_non_trivial_variable_declaration():
@@ -941,6 +942,8 @@ def test_box_and_timings():
     assert prog.to_qasm() == expected
     # Todo: box only currently technically allows QuantumStatements (i.e. gates)
     _check_respects_type_hints(prog, ["body"])
+    assert expr_matches(constant(100e-9, 0.5).args["length"], convert_float_to_duration(100e-9))
+    assert expr_matches(constant(100e-9, 0.5).args["iq"], 0.5)
 
 
 def test_play_capture():
@@ -1037,6 +1040,9 @@ def test_declare_extern():
     ).strip()
 
     assert program.to_qasm() == expected
+    assert "x" in arctan(f, f).args
+    assert expr_matches(arctan(f, i).args["x"], f)
+    assert expr_matches(arctan(f, i).args["y"], i)
 
 
 def test_defcals():