diff --git a/opensquirrel/circuit_builder.py b/opensquirrel/circuit_builder.py index 47272c18..c7f32b61 100644 --- a/opensquirrel/circuit_builder.py +++ b/opensquirrel/circuit_builder.py @@ -80,6 +80,13 @@ def __getattr__(self, attr: str) -> Any: # Default behaviour return self.__getattribute__(attr) + def __contains__(self, item: object) -> bool: + if isinstance(item, QubitRegister): + return item.name in self.register_manager.qubit_registry + if isinstance(item, BitRegister): + return item.name in self.register_manager.bit_registry + return False + def add_register(self, register: QubitRegister | BitRegister) -> None: """Add a (qu)bit register to the circuit builder. diff --git a/opensquirrel/register_manager.py b/opensquirrel/register_manager.py index fc1bd9a5..591e7e21 100644 --- a/opensquirrel/register_manager.py +++ b/opensquirrel/register_manager.py @@ -105,6 +105,14 @@ def __init__(self, qubit_registry: QubitRegistry, bit_registry: BitRegistry) -> BitRegister(0) if not bit_registry else (RegisterManager.generate_virtual_register(bit_registry)) ) + @property + def qubit_registry(self) -> QubitRegistry: + return self._qubit_registry + + @property + def bit_registry(self) -> BitRegistry: + return self._bit_registry + @property def qubit_register_size(self) -> int: return self._virtual_qubit_register.size diff --git a/tests/test_circuit_builder.py b/tests/test_circuit_builder.py index 3579267f..98ed986a 100644 --- a/tests/test_circuit_builder.py +++ b/tests/test_circuit_builder.py @@ -629,3 +629,32 @@ def test_add_instruction_non_unitaries(self, instruction: Instruction, expected: circuit = builder.to_circuit() assert circuit.ir.statements == [expected] + + +@pytest.fixture +def qubit_register1() -> QubitRegister: + return QubitRegister(4, "q0") + + +@pytest.fixture +def qubit_register2() -> QubitRegister: + return QubitRegister(8, "q1") + + +@pytest.fixture +def bit_register() -> BitRegister: + return BitRegister(4, "b0") + + +def test_qubit_register_in_builder(qubit_register1: QubitRegister, qubit_register2: QubitRegister) -> None: + builder = CircuitBuilder() + builder.add_register(qubit_register1) + assert qubit_register1 in builder + assert qubit_register2 not in builder + + +def test_bit_register_in_builder(bit_register: BitRegister) -> None: + builder = CircuitBuilder() + builder.add_register(bit_register) + assert bit_register in builder + assert BitRegister(8, "b1") not in builder