Create oqpy.Range to allow ForIn ranges with expressions. (#47)

* Add oqpy.Range

* Eliminate unnecessary int | AstConvertible

Author: PhilReinhold

oqpy/control_flow.py

@@ -38,7 +38,7 @@
     from oqpy.program import Program
 
 
-__all__ = ["If", "Else", "ForIn", "While"]
+__all__ = ["If", "Else", "ForIn", "While", "Range"]
 
 
 @contextlib.contextmanager
@@ -129,18 +129,40 @@ def ForIn(
             iterator = (make_duration(i) for i in iterator)
 
         set_declaration = ast.DiscreteSet([to_ast(program, i) for i in iterator])
-    elif isinstance(iterator, _ClassicalVar):
-        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")
+        set_declaration = to_ast(program, iterator)
 
     stmt = ast.ForInLoop(
         identifier_type.type_cls(), var.to_ast(program), set_declaration, state.body
     )
     program._add_statement(stmt)
 
 
+class Range:
+    """AstConvertible which creates an integer range.
+
+    Unlike builtin python range, this allows the components to be AstConvertible,
+    instead of just int.
+    """
+
+    def __init__(self, start: AstConvertible, stop: AstConvertible, step: AstConvertible = 1):
+        self.start = start
+        self.stop = stop
+        self.step = step
+
+    def to_ast(self, program: Program) -> ast.Expression:
+        """Convert to an ast.RangeDefinition."""
+        return ast.RangeDefinition(
+            to_ast(program, self.start),
+            ast.BinaryExpression(
+                lhs=to_ast(program, self.stop),
+                op=ast.BinaryOperator["-"],
+                rhs=ast.IntegerLiteral(value=1),
+            ),
+            to_ast(program, self.step) if self.step != 1 else None,
+        )
+
+
 @contextlib.contextmanager
 def While(program: Program, condition: OQPyExpression) -> Iterator[None]:
     """Context manager for looping a repeating a portion of a program while a condition is True.

tests/test_directives.py

@@ -1825,3 +1825,23 @@ def test_constant_conversion():
         """
     ).strip()
     assert prog.to_qasm() == expected
+
+
+def test_oqpy_range():
+    prog = Program()
+    sum = oqpy.IntVar(0, "sum")
+    with ForIn(prog, range(10), "i") as i:
+        with ForIn(prog, oqpy.Range(1, i), "j") as j:
+            prog.increment(sum, j)
+    expected = textwrap.dedent(
+        """
+        OPENQASM 3.0;
+        int[32] sum = 0;
+        for int i in [0:9] {
+            for int j in [1:i - 1] {
+                sum += j;
+            }
+        }
+        """
+    ).strip()
+    assert prog.to_qasm() == expected