Allow custom ExternArgument in extern declaration (#83)

* Allow custom ExternArgument in extern declaration

This is useful if people want to specify ArrayReferenceType as an argument to an
extern function. ArrayReferenceType require an explicit Access
modifier (mutable or readonly). Since this is just one edge case, we can let
users still specify just the type for all other extern declarations.

* Hide direct AST node manipulation

* Coverage

* Docstrings

* Rename

* Allow None as argument access

Author: anuragm

oqpy/classical_types.py

@@ -86,6 +86,7 @@
     "complex128",
     "angle_",
     "angle32",
+    "arrayreference_",
 ]
 
 # The following methods and constants are useful for creating signatures
@@ -129,6 +130,26 @@ def bit_(size: int | None = None) -> ast.BitType:
     return ast.BitType(ast.IntegerLiteral(size) if size is not None else None)
 
 
+def arrayreference_(
+    dtype: Union[
+        ast.IntType,
+        ast.UintType,
+        ast.FloatType,
+        ast.AngleType,
+        ast.DurationType,
+        ast.BitType,
+        ast.BoolType,
+        ast.ComplexType,
+    ],
+    dims: int | list[int],
+) -> ast.ArrayReferenceType:
+    """Create an array reference type."""
+    dim = (
+        ast.IntegerLiteral(dims) if isinstance(dims, int) else [ast.IntegerLiteral(d) for d in dims]
+    )
+    return ast.ArrayReferenceType(base_type=dtype, dimensions=dim)
+
+
 duration = ast.DurationType()
 stretch = ast.StretchType()
 bool_ = ast.BoolType()

oqpy/subroutines.py

@@ -19,7 +19,8 @@
 
 import functools
 import inspect
-from typing import Any, Callable, Optional, Sequence, TypeVar, get_type_hints
+from dataclasses import dataclass
+from typing import Any, Callable, Literal, Optional, Sequence, TypeVar, get_type_hints
 
 from mypy_extensions import VarArg
 from openpulse import ast
@@ -30,13 +31,26 @@
 from oqpy.quantum_types import Qubit
 from oqpy.timing import convert_float_to_duration
 
-__all__ = ["subroutine", "declare_extern", "declare_waveform_generator"]
+__all__ = ["subroutine", "declare_extern", "declare_waveform_generator", "OQPyArgument"]
 
 SubroutineParams = [oqpy.Program, VarArg(AstConvertible)]
 
 FnType = TypeVar("FnType", bound=Callable[..., Any])
 
 
+@dataclass
+class OQPyArgument:
+    """An oqpy argument to extern declaration.."""
+
+    name: str
+    dtype: ast.ClassicalType
+    access: Literal["readonly", "mutable"] | None = None
+
+    def unzip(self) -> tuple[str, ast.ClassicalType, ast.AccessControl | None]:
+        """Returns the three values, name, dtype and access as a tuple."""
+        return self.name, self.dtype, ast.AccessControl[self.access] if self.access else None
+
+
 def enable_decorator_arguments(f: FnType) -> Callable[..., FnType]:
     @functools.wraps(f)
     def decorator(*args, **kwargs):  # type: ignore[no-untyped-def]
@@ -173,7 +187,7 @@ def wrapper(
 
 def declare_extern(
     name: str,
-    args: list[tuple[str, ast.ClassicalType]],
+    args: list[tuple[str, ast.ClassicalType] | OQPyArgument],
     return_type: Optional[ast.ClassicalType] = None,
     annotations: Sequence[str | tuple[str, str]] = (),
 ) -> Callable[..., OQFunctionCall]:
@@ -190,11 +204,28 @@ def declare_extern(
         program.set(var, sqrt(0.5))
 
     """
-    arg_names = list(zip(*(args)))[0] if args else []
-    arg_types = list(zip(*(args)))[1] if args else []
+    arg_names: list[str] = []
+    arg_types: list[ast.ClassicalType] = []
+    arg_access: list[ast.AccessControl | None] = []
+
+    for arg in args:
+        if isinstance(arg, tuple):
+            arg_name, arg_type = arg
+            access = None
+        elif isinstance(arg, OQPyArgument):
+            arg_name, arg_type, access = arg.unzip()
+        else:
+            raise Exception(f"Argument {arg} should have a proper type")
+        arg_names.append(arg_name)
+        arg_types.append(arg_type)
+        arg_access.append(access)
+
     extern_decl = ast.ExternDeclaration(
         ast.Identifier(name),
-        [ast.ExternArgument(type=t) for t in arg_types],
+        [
+            ast.ExternArgument(type=ctype, access=access)
+            for ctype, access in zip(arg_types, arg_access)
+        ],
         return_type,
     )
     extern_decl.annotations = make_annotations(annotations)
@@ -236,7 +267,7 @@ def call_extern(*call_args: AstConvertible, **call_kwargs: AstConvertible) -> OQ
 
 def declare_waveform_generator(
     name: str,
-    argtypes: list[tuple[str, ast.ClassicalType]],
+    argtypes: list[tuple[str, ast.ClassicalType] | OQPyArgument],
     annotations: Sequence[str | tuple[str, str]] = (),
 ) -> Callable[..., OQFunctionCall]:
     """Create a function which generates waveforms using a specified name and argument signature."""

tests/test_directives.py

@@ -1024,6 +1024,19 @@ def test_declare_extern():
     # Test an extern with no input and no output
     fire_bazooka = declare_extern("fire_bazooka", [])
 
+    # Test an extern with readonly array
+    print_array = declare_extern(
+        "print_array",
+        [
+            ("style", int32),
+            OQPyArgument(
+                name="arr",
+                dtype=arrayreference_(int32, 1),
+                access="readonly",
+            ),
+        ],
+    )
+
     f = oqpy.FloatVar(name="f", init_expression=0.0)
     i = oqpy.IntVar(name="i", init_expression=5)
 
@@ -1032,6 +1045,7 @@ def test_declare_extern():
     program.set(i, time())
     program.do_expression(set_global_voltage(i))
     program.do_expression(fire_bazooka())
+    program.do_expression(print_array(1, [0, 1, 2]))
 
     expected = textwrap.dedent(
         """
@@ -1041,13 +1055,15 @@ def test_declare_extern():
         extern time() -> int[32];
         extern set_voltage(int[32]);
         extern fire_bazooka();
+        extern print_array(int[32], readonly array[int[32], #dim=1]);
         float[64] f = 0.0;
         int[32] i = 5;
         f = sqrt(f);
         f = arctan(f, f);
         i = time();
         set_voltage(i);
         fire_bazooka();
+        print_array(1, {0, 1, 2});
         """
     ).strip()
 
@@ -1057,6 +1073,11 @@ def test_declare_extern():
     assert expr_matches(arctan(f, i).args["y"], i)
 
 
+def test_invalid_extern_declaration():
+    # Test with invalid argument
+    with pytest.raises(Exception, match="Argument.*"):
+        _ = declare_extern("invalid", [int32])
+
 def test_defcals():
     prog = Program()
     constant = declare_waveform_generator("constant", [("length", duration), ("iq", complex128)])