Allow custom ExternArgument in extern declaration

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