Finalize Python Registers, prepare for replacements.

raynelfss · Jan 3, 2025 · f380865 · f380865
1 parent f80c89d
commit f380865
Show file tree

Hide file tree

Showing 3 changed files with 397 additions and 23 deletions.
diff --git a/crates/circuit/src/bit.rs b/crates/circuit/src/bit.rs
@@ -1,6 +1,6 @@
 use std::hash::{DefaultHasher, Hash, Hasher};
 
-use pyo3::{prelude::*, types::PyDict};
+use pyo3::{exceptions::PyTypeError, prelude::*, types::PyDict};
 
 use crate::{
     circuit_data::CircuitError,
@@ -90,13 +90,33 @@ impl PyBit {
     fn __setstate__(mut slf: PyRefMut<'_, Self>, state: (Option<(String, u32)>, Option<u32>)) {
         slf.register = state
             .0
-            .map(|(name, num_qubits)| RegisterAsKey::new(name.as_str(), num_qubits));
+            .map(|(name, num_qubits)| RegisterAsKey::Register((name, num_qubits)));
         slf.index = state.1;
     }
+
+    fn __repr__(slf: Bound<Self>) -> PyResult<String> {
+        let borrowed = slf.borrow();
+        if borrowed.register.is_none() && borrowed.index.is_none() {
+            return Ok(slf.py_super()?.repr()?.to_string());
+        }
+        let reg = borrowed.register.as_ref().unwrap();
+        Ok(format!(
+            "{}({}({:?}, {}), {})",
+            slf.get_type().name()?,
+            reg.type_identifier(),
+            reg.name(),
+            reg.index(),
+            borrowed.index.unwrap()
+        ))
+    }
+
+    pub fn is_new(&self) -> bool {
+        self.index.is_none() && self.register.is_none()
+    }
 }
 
 macro_rules! create_py_bit {
-    ($name:ident, $pyname:literal, $module:literal) => {
+    ($name:ident, $pyname:literal, $reg_type:pat, $module:literal) => {
         #[pyclass(name=$pyname, extends=PyBit, subclass, module=$module)]
         #[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
         pub struct $name();
@@ -109,15 +129,23 @@ macro_rules! create_py_bit {
                 register: Option<RegisterAsKey>,
                 index: Option<u32>,
             ) -> PyResult<(Self, PyBit)> {
-                Ok((Self(), PyBit::new(register, index)?))
+                if register.is_none() || matches!(register, Some($reg_type)) {
+                    Ok((Self(), PyBit::new(register, index)?))
+                } else {
+                    Err(PyTypeError::new_err(format!(
+                        "The incorrect register was assigned. Bit type {}, Register type {}",
+                        $pyname,
+                        register.unwrap().type_identifier()
+                    )))
+                }
             }
         }
     };
 }
 
 // Create python instances
-create_py_bit! {PyQubit, "Qubit", "qiskit._accelerate.bit"}
-create_py_bit! {PyClbit, "Clbit", "qiskit._accelerate.bit"}
+create_py_bit! {PyQubit, "Qubit", RegisterAsKey::Quantum(_), "qiskit._accelerate.bit"}
+create_py_bit! {PyClbit, "Clbit", RegisterAsKey::Classical(_), "qiskit._accelerate.bit"}
 
 /// Keeps information about where a qubit is located within the circuit.
 #[derive(Debug, Clone)]

diff --git a/crates/circuit/src/lib.rs b/crates/circuit/src/lib.rs
@@ -128,6 +128,9 @@ pub fn circuit(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_class::<bit::PyBit>()?;
     m.add_class::<bit::PyClbit>()?;
     m.add_class::<bit::PyQubit>()?;
+    m.add_class::<register::PyRegister>()?;
+    m.add_class::<register::PyQuantumRegister>()?;
+    m.add_class::<register::PyClassicalRegister>()?;
     m.add_class::<circuit_data::CircuitData>()?;
     m.add_class::<circuit_instruction::CircuitInstruction>()?;
     m.add_class::<dag_circuit::DAGCircuit>()?;