fix: refactor qubit types into new submodule

src/autoqasm/__init__.py

@@ -50,12 +50,13 @@ def my_program():
 from .api import gate, gate_calibration, main, subroutine  # noqa: F401
 from .instructions import QubitIdentifierType as Qubit  # noqa: F401
 from .program import Program, build_program, verbatim  # noqa: F401
+from .program import get_program_conversion_context as _get_program_conversion_context
 from .transpiler import transpiler  # noqa: F401
 from .types import ArrayVar, BitVar, BoolVar, FloatVar, IntVar  # noqa: F401
 from .types import Range as range  # noqa: F401
 
 
 def __getattr__(name):
     if name == "qubits":
-        return instructions.global_qubit_register()
+        return _get_program_conversion_context().global_qubit_register
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/autoqasm/instructions/__init__.py

@@ -28,4 +28,3 @@ def bell():
 from .gates import *
 from .instructions import barrier, reset  # noqa: F401
 from .measurements import measure, measure_ff  # noqa: F401
-from .qubits import global_qubit_register  # noqa: F401

src/autoqasm/instructions/instructions.py

@@ -23,8 +23,9 @@
 
 from autoqasm import program as aq_program
 from autoqasm import types as aq_types
-from autoqasm.instructions.qubits import _as_qubit_iterable, _qubit
+from autoqasm.instructions.qubits import _qubit
 from autoqasm.types import QubitIdentifierType
+from autoqasm.types.qubits import _as_qubit_iterable
 from braket.circuits.basis_state import BasisState, BasisStateInput
 
 

src/autoqasm/instructions/measurements.py

@@ -29,7 +29,8 @@ def my_program():
 from autoqasm import program
 from autoqasm import types as aq_types
 from autoqasm.instructions.instructions import _qubit_instruction
-from autoqasm.instructions.qubits import _as_qubit_iterable, _qubit, global_qubit_register
+from autoqasm.instructions.qubits import _qubit
+from autoqasm.types.qubits import _as_qubit_iterable
 
 
 def measure(
@@ -45,9 +46,10 @@ def measure(
     Returns:
         BitVar: Bit variable the measurement results are assigned to.
     """
-    qubits = _as_qubit_iterable(qubits, default=global_qubit_register())
+    ctx = program.get_program_conversion_context()
+    qubits = _as_qubit_iterable(qubits, default=ctx.global_qubit_register)
 
-    oqpy_program = program.get_program_conversion_context().get_oqpy_program()
+    oqpy_program = ctx.get_oqpy_program()
 
     bit_var_size = len(qubits) if len(qubits) > 1 else None
     bit_var = aq_types.BitVar(

src/autoqasm/instructions/qubits.py

@@ -17,15 +17,13 @@
 from __future__ import annotations
 
 import re
-from collections.abc import Iterable, Iterator
 from functools import singledispatch
 from typing import Any
 
 import oqpy.base
 from openpulse.printer import dumps
 
-from autoqasm import constants, errors, program
-from autoqasm import types as aq_types
+from autoqasm import errors, program
 
 
 def _get_physical_qubit_indices(qids: list[str]) -> list[int]:
@@ -48,39 +46,6 @@ def _get_physical_qubit_indices(qids: list[str]) -> list[int]:
     return braket_qubits
 
 
-def _global_qubit_register(qubit_idx_expr: int | str) -> oqpy.Qubit:
-    return oqpy.Qubit(f"{constants.QUBIT_REGISTER}[{qubit_idx_expr}]", needs_declaration=False)
-
-
-class GlobalQubitRegister(oqpy.Qubit):
-    def __init__(self, size: int | None):
-        super().__init__(name=constants.QUBIT_REGISTER, size=size, needs_declaration=False)
-
-    def __len__(self) -> int:
-        return self.size
-
-    def __iter__(self) -> Iterator:
-        return iter(range(len(self)))
-
-
-def global_qubit_register() -> GlobalQubitRegister:
-    return program.get_program_conversion_context().global_qubit_register
-
-
-def _as_qubit_iterable(
-    qubits: aq_types.QubitIdentifierType | Iterable[aq_types.QubitIdentifierType] | None,
-    default: Iterable[aq_types.QubitIdentifierType] | None = None,
-) -> Iterable[aq_types.QubitIdentifierType]:
-    """Normalize a qubit argument to an iterable. ``None`` maps to ``default`` (an empty
-    list if not provided); a single qubit identifier is wrapped in a list; iterables
-    (including :class:`GlobalQubitRegister`) pass through unchanged."""
-    if qubits is None:
-        qubits = default if default is not None else []
-    if aq_types.is_qubit_identifier_type(qubits) and not isinstance(qubits, GlobalQubitRegister):
-        return [qubits]
-    return qubits
-
-
 @singledispatch
 def _qubit(qid: Any) -> oqpy.Qubit:
     """Maps a given qubit representation to an oqpy qubit.
@@ -91,48 +56,40 @@ def _qubit(qid: Any) -> oqpy.Qubit:
     Returns:
         Qubit: A translated oqpy qubit.
     """
-    raise ValueError(f"invalid qubit label: '{qid}'")
+    raise TypeError(f"object of type '{type(qid).__name__}' cannot be used as a qubit")
 
 
 @_qubit.register
 def _(qid: bool) -> oqpy.Qubit:
-    raise ValueError(f"invalid qubit label: '{qid}'")
-
-
-@_qubit.register
-def _(qid: float) -> oqpy.Qubit:
-    raise TypeError(f"qubit index cannot be a float: '{qid}'")
+    raise TypeError(f"object of type '{type(qid).__name__}' cannot be used as a qubit")
 
 
 @_qubit.register
 def _(qid: int) -> oqpy.Qubit:
     # Integer virtual qubit, like `h(0)`
-    program.get_program_conversion_context().register_qubit(qid)
-    return _global_qubit_register(qid)
-
-
-@_qubit.register
-def _(qid: GlobalQubitRegister) -> oqpy.Qubit:
-    raise ValueError("qubit index must be a single value, not a list or a register")
+    ctx = program.get_program_conversion_context()
+    ctx.register_qubit(qid)
+    return ctx.global_qubit_register[qid]
 
 
 @_qubit.register
 def _(qid: oqpy._ClassicalVar) -> oqpy.Qubit:
     # Indexed by variable, such as i in range(n); h(i)
-    if program.get_program_conversion_context().get_declared_qubits() is None:
+    ctx = program.get_program_conversion_context()
+    if ctx.get_declared_qubits() is None:
         raise errors.UnknownQubitCountError()
-    return _global_qubit_register(qid.name)
+    return ctx.global_qubit_register[qid.name]
 
 
 @_qubit.register
 def _(qid: oqpy.base.OQPyExpression) -> oqpy.Qubit:
     # Indexed by expression, such as i in range(n); h(i + 1)
-    if program.get_program_conversion_context().get_declared_qubits() is None:
+    ctx = program.get_program_conversion_context()
+    if ctx.get_declared_qubits() is None:
         raise errors.UnknownQubitCountError()
 
-    oqpy_program = program.get_program_conversion_context().get_oqpy_program()
-    qubit_idx_expr = dumps(qid.to_ast(oqpy_program))
-    return _global_qubit_register(qubit_idx_expr)
+    qubit_idx_expr = dumps(qid.to_ast(ctx.get_oqpy_program()))
+    return ctx.global_qubit_register[qubit_idx_expr]
 
 
 @_qubit.register

src/autoqasm/operators/control_flow.py

@@ -22,11 +22,11 @@
 import oqpy.base
 
 from autoqasm import program
-from autoqasm.types import Range, is_qasm_type
+from autoqasm.types import GlobalQubitRegister, Range, is_qasm_type
 
 
 def for_stmt(
-    iter: Iterable | oqpy.Range | oqpy.Qubit,
+    iter: Iterable | oqpy.Range,
     extra_test: Callable[[], Any] | None,
     body: Callable[[Any], None],
     get_state: Any,
@@ -37,7 +37,7 @@ def for_stmt(
     """Implements a for loop.
 
     Args:
-        iter (Iterable | Range | Qubit): The iterable to be looped over.
+        iter (Iterable | Range): The iterable to be looped over.
         extra_test (Callable[[], Any] | None): A function to cause the loop to break if true.
         body (Callable[[Any], None]): The body of the for loop.
         get_state (Any): Unused.
@@ -56,13 +56,13 @@ def for_stmt(
 
 
 def _oqpy_for_stmt(
-    iter: oqpy.Range | oqpy.Qubit,
+    iter: oqpy.Range | GlobalQubitRegister,
     body: Callable[[Any], None],
     opts: dict,
 ) -> None:
     """Overload of for_stmt that produces an oqpy for loop."""
     ctx = program.get_program_conversion_context()
-    if isinstance(iter, oqpy.Qubit):
+    if isinstance(iter, GlobalQubitRegister):
         iter = Range(iter.size)
 
     def _trace(ctx):

src/autoqasm/program/program.py

@@ -34,11 +34,12 @@
 
 import autoqasm.types as aq_types
 from autoqasm import _frame_filtering, constants, errors
-from autoqasm.instructions.qubits import GlobalQubitRegister, _get_physical_qubit_indices, _qubit
+from autoqasm.instructions.qubits import _get_physical_qubit_indices, _qubit
 from autoqasm.program.serialization_properties import (
     OpenQASMSerializationProperties,
     SerializationProperties,
 )
+from autoqasm.types import GlobalQubitRegister
 from autoqasm.types import QubitIdentifierType as Qubit
 from autoqasm.types.deferred import DeferredVarMixin
 from braket.aws.aws_device import AwsDevice
@@ -422,7 +423,7 @@ def declare_global_qubit_register(self, size: int) -> None:
         """
         root_oqpy_program = self.get_oqpy_program(scope=ProgramScope.MAIN)
         self.global_qubit_register.size = size
-        root_oqpy_program.declare(self.global_qubit_register, to_beginning=True)
+        root_oqpy_program.declare(self.global_qubit_register.oqpy_var, to_beginning=True)
 
     def register_gate(self, gate_name: str, is_compiler_directive: bool = False) -> None:
         """Register a gate that is used in this program.
@@ -650,17 +651,28 @@ def is_var_name_used(self, var_name: str) -> bool:
         oqpy_program = self.get_oqpy_program()
         return var_name in oqpy_program.declared_vars or var_name in oqpy_program.undeclared_vars
 
-    def validate_gate_targets(self, qubits: list[Any], angles: list[Any]) -> None:
+    def validate_gate_targets(
+        self,
+        qubits: Iterable[Qubit],
+        angles: Iterable[Any],
+    ) -> None:
         """Validate that the specified gate targets are valid at this point in the program.
 
         Args:
-            qubits (list[Any]): The list of target qubits to validate.
-            angles (list[Any]): The list of target angles to validate.
+            qubits (Iterable[QubitIdentifierType]): The target qubits to validate.
+            angles (Iterable[Any]): The target angles to validate.
 
         Raises:
+            TypeError: A target qubit is not a single qubit identifier.
             errors.InvalidTargetQubit: Target qubits are invalid in the current context.
             errors.InvalidGateDefinition: Targets are invalid in the current gate definition.
         """
+        for qubit in qubits:
+            if not aq_types.is_qubit_identifier_type(qubit):
+                raise TypeError(
+                    f'Invalid qubit target: "{qubit}". Target must be a single qubit, not a list or a register.'
+                )
+
         if self.in_verbatim_block and not self._gate_definitions_processing:
             self._validate_verbatim_target_qubits(qubits)
 

src/autoqasm/types/__init__.py

@@ -16,15 +16,18 @@
 """
 
 from .conversions import map_parameter_type, var_type_from_oqpy, wrap_value  # noqa: F401
+from .qubits import (  # noqa: F401
+    GlobalQubitRegister,
+    QubitIdentifierType,
+    is_qubit_identifier_type,
+)
 from .types import (  # noqa: F401
     ArrayVar,
     BitVar,
     BoolVar,
     FloatVar,
     IntVar,
-    QubitIdentifierType,
     Range,
     is_qasm_type,
-    is_qubit_identifier_type,
     make_annotations_list,
 )

src/autoqasm/types/qubits.py

@@ -0,0 +1,90 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+"""Qubit identifier types and the program-level qubit register."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from typing import Any, get_args
+
+import oqpy
+import oqpy.base
+
+from autoqasm import constants
+from braket.registers import Qubit
+
+QubitIdentifierType = int | str | Qubit | oqpy._ClassicalVar | oqpy.base.OQPyExpression | oqpy.Qubit
+
+_QUBIT_IDENTIFIER_TYPES: tuple[type, ...] = get_args(QubitIdentifierType)
+
+
+def is_qubit_identifier_type(qubit: Any) -> bool:
+    """Checks if a given object is a qubit identifier type.
+
+    Args:
+        qubit (Any): The object to check.
+
+    Returns:
+        bool: True if the object is a qubit identifier type, False otherwise.
+    """
+    return isinstance(qubit, _QUBIT_IDENTIFIER_TYPES)
+
+
+def _as_qubit_iterable(
+    qubits: QubitIdentifierType | Iterable[QubitIdentifierType] | None,
+    default: Iterable[QubitIdentifierType] | None = None,
+) -> Iterable[QubitIdentifierType]:
+    """Normalize a qubit argument to an iterable. ``None`` maps to ``default`` (an empty
+    list if not provided); a single qubit identifier is wrapped in a list; iterables
+    (including :class:`GlobalQubitRegister`) pass through unchanged."""
+    if qubits is None:
+        qubits = default if default is not None else []
+    if is_qubit_identifier_type(qubits):
+        return [qubits]
+    return qubits
+
+
+class GlobalQubitRegister:
+    def __init__(self, size: int | None = None):
+        self._var = oqpy.Qubit(constants.QUBIT_REGISTER, size=size, needs_declaration=False)
+
+    @property
+    def name(self) -> str:
+        return self._var.name
+
+    @property
+    def size(self) -> int | None:
+        return self._var.size
+
+    @size.setter
+    def size(self, value: int) -> None:
+        self._var.size = value
+
+    @property
+    def oqpy_var(self) -> oqpy.Qubit:
+        """The underlying oqpy variable used to declare the register in OpenQASM."""
+        return self._var
+
+    def __repr__(self) -> str:
+        return self.name
+
+    def __len__(self) -> int:
+        return self.size
+
+    def __iter__(self) -> Iterator[int]:
+        return iter(range(len(self)))
+
+    def __getitem__(self, index: int | str) -> oqpy.Qubit:
+        """Returns an oqpy.Qubit referring to ``__qubits__[index]``."""
+        return oqpy.Qubit(f"{self.name}[{index}]", needs_declaration=False)