Allow for loops with arbitrary variable type

oqpy/control_flow.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Iterable, Iterator, Optional
+from typing import TYPE_CHECKING, Iterable, Iterator, Optional, Type
 
 from openpulse import ast
 
@@ -73,7 +73,8 @@ def ForIn(
     program: Program,
     iterator: Iterable[AstConvertible] | range | AstConvertible,
     identifier_name: Optional[str] = None,
-) -> Iterator[IntVar]:
+    identifier_type: Type[_ClassicalVar] = IntVar,
+) -> Iterator[_ClassicalVar]:
     """Context manager for looping a particular portion of a program.
 
     .. code-block:: python
@@ -84,20 +85,23 @@ def ForIn(
 
     """
     program._push()
-    var = IntVar(name=identifier_name, needs_declaration=False)
+    var = identifier_type(name=identifier_name, needs_declaration=False)
     yield var
     state = program._pop()
 
     if isinstance(iterator, range):
-        iterator = convert_range(program, iterator)
+        set_declaration = convert_range(program, iterator)
     elif isinstance(iterator, Iterable):
-        iterator = ast.DiscreteSet([to_ast(program, i) for i in iterator])
+        set_declaration = ast.DiscreteSet([to_ast(program, i) for i in iterator])
     elif isinstance(iterator, _ClassicalVar):
-        iterator = to_ast(program, iterator)
+        set_declaration = to_ast(program, iterator)
+        assert isinstance(set_declaration, ast.Identifier), type(set_declaration)
     else:
         raise TypeError(f"'{type(iterator)}' object is not iterable")
 
-    stmt = ast.ForInLoop(ast.IntType(size=None), var.to_ast(program), iterator, state.body)
+    stmt = ast.ForInLoop(
+        identifier_type.type_cls(), var.to_ast(program), set_declaration, state.body
+    )
     program._add_statement(stmt)
 
 

tests/test_directives.py

@@ -20,6 +20,7 @@
 
 import numpy as np
 import pytest
+from openpulse import ast
 from openpulse.printer import dumps
 
 import oqpy
@@ -330,6 +331,54 @@ def test_for_in():
     assert prog.to_qasm() == expected
 
 
+def test_for_in_var_types():
+    port = oqpy.PortVar("my_port")
+    frame = oqpy.FrameVar(port, 3e9, 0, "my_frame")
+
+    # Test over floating point array.
+    program = oqpy.Program()
+    frequencies = [0.1, 0.2, 0.5]
+    with oqpy.ForIn(program, frequencies, "frequency", oqpy.FloatVar) as f:
+        program.set_frequency(frame, f)
+
+    expected = textwrap.dedent(
+        """
+        OPENQASM 3.0;
+        port my_port;
+        frame my_frame = newframe(my_port, 3000000000.0, 0);
+        for float frequency in {0.1, 0.2, 0.5} {
+            set_frequency(my_frame, frequency);
+        }
+        """
+    ).strip()
+
+    assert program.to_qasm() == expected
+
+    # Test over duration array.
+    program = oqpy.Program()
+    float_delays = [1e-9, 2e-9, 5e-9, 10e-9, 1e-6]
+    duration_delays = [
+        ast.DurationLiteral(round(1e9 * float_delay), ast.TimeUnit.ns)
+        for float_delay in float_delays
+    ]
+
+    with oqpy.ForIn(program, duration_delays, "d", DurationVar) as delay:
+        program.delay(delay, frame)
+
+    expected = textwrap.dedent(
+        """
+        OPENQASM 3.0;
+        port my_port;
+        frame my_frame = newframe(my_port, 3000000000.0, 0);
+        for duration d in {1ns, 2ns, 5ns, 10ns, 1000ns} {
+            delay[d] my_frame;
+        }
+        """
+    ).strip()
+
+    assert program.to_qasm() == expected
+
+
 def test_while():
     prog = Program()
     j = IntVar(0, "j")