Qubit and BitVar supports get item (#67)

* Support BitVar indexing

* move validation to OQIndexExpression

* qubit class supports get item

* add IndexedQubitArray class

* move OQIndexExpression to classical_types

* Update oqpy/quantum_types.py

* Update oqpy/classical_types.py

* Update oqpy/classical_types.py

* Apply suggestions from code review

---------

Co-authored-by: Phil Reinhold <[email protected]>

Author: yitchen-tim

oqpy/classical_types.py

@@ -62,7 +62,6 @@
     "ComplexVar",
     "DurationVar",
     "OQFunctionCall",
-    "OQIndexExpression",
     "StretchVar",
     "_ClassicalVar",
     "duration",
@@ -236,6 +235,25 @@ def __init__(self, *args: Any, size: int | None = None, **kwargs: Any):
             self.size = size
         super().__init__(*args, **kwargs, size=ast.IntegerLiteral(self.size) if self.size else None)
 
+    def _validate_getitem_index(self, index: AstConvertible) -> None:
+        """Validate the index and variable for `__getitem__`.
+
+        Args:
+            var (_SizedVar): Variable to apply `__getitem__`.
+            index (AstConvertible): Index for `__getitem__`.
+        """
+        if self.size is None:
+            raise TypeError(f"'{self.name}' is not subscriptable")
+
+        if isinstance(index, int):
+            if not 0 <= index < self.size:
+                raise IndexError("list index out of range.")
+        elif isinstance(index, OQPyExpression):
+            if not isinstance(index.type, (ast.IntType, ast.UintType)):
+                raise IndexError("The list index must be an integer.")
+        else:
+            raise IndexError("The list index must be an integer.")
+
 
 _SizedVarT = TypeVar("_SizedVarT", bound=_SizedVar)
 
@@ -273,22 +291,9 @@ class BitVar(_SizedVar):
 
     type_cls = ast.BitType
 
-    def __getitem__(self, idx: Union[int, slice, Iterable[int]]) -> BitVar:
-        if self.size is None:
-            raise TypeError(f"'{self.type_cls}' object is not subscriptable")
-        if isinstance(idx, int):
-            if 0 <= idx < self.size:
-                return BitVar(
-                    init_expression=ast.IndexExpression(
-                        ast.Identifier(self.name), [ast.IntegerLiteral(idx)]
-                    ),
-                    name=f"{self.name}[{idx}]",
-                    needs_declaration=False,
-                )
-            else:
-                raise IndexError("list index out of range.")
-        else:
-            raise IndexError("The list index must be an integer.")
+    def __getitem__(self, index: AstConvertible) -> OQIndexExpression:
+        self._validate_getitem_index(index)
+        return OQIndexExpression(collection=self, index=index, type_=self.type_cls())
 
 
 class ComplexVar(_ClassicalVar):
@@ -394,18 +399,16 @@ def __init__(
         )
 
     def __getitem__(self, index: AstConvertible) -> OQIndexExpression:
-        return OQIndexExpression(collection=self, index=index)
+        return OQIndexExpression(collection=self, index=index, type_=self.base_type().type_cls())
 
 
 class OQIndexExpression(OQPyExpression):
     """An oqpy expression corresponding to an index expression."""
 
-    def __init__(self, collection: AstConvertible, index: AstConvertible):
+    def __init__(self, collection: AstConvertible, index: AstConvertible, type_: ast.ClassicalType):
         self.collection = collection
         self.index = index
-
-        if isinstance(collection, ArrayVar):
-            self.type = collection.base_type().type_cls()
+        self.type = type_
 
     def to_ast(self, program: Program) -> ast.IndexExpression:
         """Converts this oqpy index expression into an ast node."""

oqpy/program.py

@@ -483,7 +483,7 @@ def gate(
         self, qubits: AstConvertible | Iterable[AstConvertible], name: str, *args: Any
     ) -> Program:
         """Apply a gate to a qubit or set of qubits."""
-        if isinstance(qubits, quantum_types.Qubit):
+        if isinstance(qubits, (quantum_types.Qubit, quantum_types.IndexedQubitArray)):
             qubits = [qubits]
         assert isinstance(qubits, Iterable)
         self._add_statement(

oqpy/quantum_types.py

@@ -30,7 +30,7 @@
     from oqpy.program import Program
 
 
-__all__ = ["Qubit", "QubitArray", "defcal", "gate", "PhysicalQubits", "Cal"]
+__all__ = ["Qubit", "defcal", "gate", "PhysicalQubits", "Cal"]
 
 
 class Qubit(Var):
@@ -39,11 +39,13 @@ class Qubit(Var):
     def __init__(
         self,
         name: str,
+        size: Optional[int] = None,
         needs_declaration: bool = True,
         annotations: Sequence[str | tuple[str, str]] = (),
     ):
         super().__init__(name, needs_declaration=needs_declaration)
         self.name = name
+        self.size = size
         self.annotations = annotations
 
     def to_ast(self, prog: Program) -> ast.Expression:
@@ -53,10 +55,18 @@ def to_ast(self, prog: Program) -> ast.Expression:
 
     def make_declaration_statement(self, program: Program) -> ast.Statement:
         """Make an ast statement that declares the OQpy variable."""
-        decl = ast.QubitDeclaration(ast.Identifier(self.name), size=None)
+        decl = ast.QubitDeclaration(
+            ast.Identifier(self.name),
+            size=ast.IntegerLiteral(self.size) if self.size else self.size,
+        )
         decl.annotations = make_annotations(self.annotations)
         return decl
 
+    def __getitem__(self, index: AstConvertible) -> IndexedQubitArray:
+        if self.size is None:
+            raise TypeError(f"'{self.name}' is not subscriptable")
+        return IndexedQubitArray(collection=self, index=index)
+
 
 class PhysicalQubits:
     """Provides a means of accessing qubit variables corresponding to physical qubits.
@@ -68,9 +78,18 @@ def __class_getitem__(cls, item: int) -> Qubit:
         return Qubit(f"${item}", needs_declaration=False)
 
 
-# Todo (#51): support QubitArray
-class QubitArray:
-    """Represents an array of qubits."""
+class IndexedQubitArray:
+    """Represents an indexed qubit array."""
+
+    def __init__(self, collection: Qubit, index: AstConvertible):
+        self.collection = collection
+        self.index = index
+
+    def to_ast(self, program: Program) -> ast.IndexExpression:
+        """Converts this indexed qubit array into an ast node."""
+        return ast.IndexExpression(
+            collection=to_ast(program, self.collection), index=[to_ast(program, self.index)]
+        )
 
 
 @contextlib.contextmanager

tests/test_directives.py

@@ -31,6 +31,7 @@
 import oqpy
 from oqpy import *
 from oqpy.base import OQPyExpression, expr_matches, logical_and, logical_or
+from oqpy.classical_types import OQIndexExpression
 from oqpy.quantum_types import PhysicalQubits
 from oqpy.timing import OQDurationLiteral
 
@@ -126,6 +127,9 @@ def test_variable_declaration():
     prog = Program(version=None)
     prog.declare(vars)
     prog.set(arr[1], 0)
+    index = IntVar(2, "index")
+    prog.set(arr[index], 1)
+    prog.set(arr[index + 1], 0)
 
     with pytest.raises(IndexError):
         prog.set(arr[40], 2)
@@ -142,6 +146,7 @@ def test_variable_declaration():
 
     expected = textwrap.dedent(
         """
+        int[32] index = 2;
         bool b = true;
         int[32] i = -4;
         uint[32] u = 5;
@@ -151,10 +156,12 @@ def test_variable_declaration():
         bit[20] arr;
         bit c;
         arr[1] = 0;
+        arr[index] = 1;
+        arr[index + 1] = 0;
         """
     ).strip()
 
-    assert isinstance(arr[14], BitVar)
+    assert isinstance(arr[14], OQIndexExpression)
     assert prog.to_qasm() == expected
     _check_respects_type_hints(prog)
 
@@ -2215,7 +2222,12 @@ def test_invalid_gates():
 def test_gate_declarations():
     prog = oqpy.Program()
     q = oqpy.Qubit("q", needs_declaration=False)
-    with oqpy.gate(prog, q, "u", [oqpy.AngleVar(name="alpha"), oqpy.AngleVar(name="beta"), oqpy.AngleVar(name="gamma")]) as (alpha, beta, gamma):
+    with oqpy.gate(
+        prog,
+        q,
+        "u",
+        [oqpy.AngleVar(name="alpha"), oqpy.AngleVar(name="beta"), oqpy.AngleVar(name="gamma")],
+    ) as (alpha, beta, gamma):
         prog.gate(q, "a", alpha)
         prog.gate(q, "b", beta)
         prog.gate(q, "c", gamma)
@@ -2266,3 +2278,25 @@ def test_include():
 
     assert prog.to_qasm() == expected
 
+
+def test_qubit_array():
+    prog = oqpy.Program()
+    q = oqpy.Qubit("q", size=2)
+    prog.gate(q[0], "h")
+    prog.gate([q[0], q[1]], "cnot")
+
+    expected = textwrap.dedent(
+        """
+        OPENQASM 3.0;
+        qubit[2] q;
+        h q[0];
+        cnot q[0], q[1];
+        """
+    ).strip()
+
+    assert prog.to_qasm() == expected
+
+    with pytest.raises(TypeError):
+        prog = oqpy.Program()
+        q = oqpy.Qubit("q")
+        prog.gate(q[0], "h")