Fix subroutines with array arguments

oqpy/subroutines.py

@@ -100,7 +100,16 @@ def increment_variable(int[32] i) {
     for argname in argnames[1:]:  # arg 0 should be program
         if argname not in type_hints:
             raise ValueError(f"No type hint provided for {argname} on subroutine {name}.")
-        elif not issubclass(type_hints[argname], (_ClassicalVar, Qubit)):
+
+        # ArrayVar[] returns a partial function instead of a type.
+        # The underlying function of that partial should be ArrayVar itself.
+        type_hint = (
+            type_hints[argname].func
+            if isinstance(type_hints[argname], functools.partial)
+            else type_hints[argname]
+        )
+
+        if not issubclass(type_hint, (_ClassicalVar, Qubit)):
             raise ValueError(
                 f"Type hint for {argname} on subroutine {name} is not an oqpy variable type."
             )

tests/test_directives.py

@@ -817,6 +817,13 @@ def delay50ns(prog: Program, q: Qubit) -> None:
     q = PhysicalQubits[0]
     prog.do_expression(delay50ns(prog, q))
 
+    @subroutine
+    def get(prog: Program, arr: ArrayVar[IntVar, 3], i: IntVar) -> IntVar:
+        return arr[i]
+
+    arr = ArrayVar[IntVar, 3]([0, 1, 2], name="arr")
+    prog.set(y, get(prog, arr, y))
+
     with pytest.raises(ValueError):
 
         @subroutine
@@ -858,9 +865,14 @@ def multiply(int[32] x, int[32] y) -> int[32] {
         def delay50ns(qubit q) {
             delay[50.0ns] q;
         }
+        def get(array[int[32], 3] arr, int[32] i) -> int {
+            return arr[i];
+        }
         int[32] y = 2;
+        array[int[32], 3] arr = {0, 1, 2};
         y = multiply(y, 3);
         delay50ns($0);
+        y = get(arr, y);
         """
     ).strip()