diff --git a/opensquirrel/circuit_builder.py b/opensquirrel/circuit_builder.py index 47272c18..0bf18ee3 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 # noqa: SLF001 + if isinstance(item, BitRegister): + return item.name in self.register_manager._bit_registry # noqa: SLF001 + return False + def add_register(self, register: QubitRegister | BitRegister) -> None: """Add a (qu)bit register to the circuit builder. 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