feat(python): Constructors, iterators, and other utility methods have been added to make working with ResultData, RegisterMap, and others easier.

crates/lib/src/execution_data.rs

@@ -83,7 +83,7 @@ pub enum RegisterMatrix {
 /// register.
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[repr(transparent)]
-pub struct RegisterMap(HashMap<String, RegisterMatrix>);
+pub struct RegisterMap(pub HashMap<String, RegisterMatrix>);
 
 /// Errors that may occur when trying to build a [`RegisterMatrix`] from execution data
 #[allow(missing_docs)]

crates/python/qcs_sdk/__init__.pyi

@@ -261,6 +261,11 @@ class ResultData:
         - ``from_*``: wrap underlying values as this enum type.
     """
 
+    def __new__(cls, inner: Union[QPUResultData, QVMResultData]) -> "ResultData":
+        """
+        Create a new ResultData from either QVM or QPU result data.
+        """
+        ...
     def to_register_map(self) -> RegisterMap:
         """
         Convert ``ResultData`` from its inner representation as ``QVMResultData`` or
@@ -299,6 +304,7 @@ class ResultData:
 
 @final
 class ExecutionData:
+    def __new__(cls, result_data: ResultData, duration: Optional[datetime.timedelta] = None): ...
     @property
     def result_data(self) -> ResultData: ...
     @result_data.setter
@@ -332,6 +338,9 @@ class RegisterData:
     ) -> Union[List[List[int]], List[List[float]], List[List[complex]]]:
         """Returns the inner value."""
         ...
+    def as_ndarray(self) -> Union[NDArray[np.int64], NDArray[np.float64], NDArray[np.complex128]]:
+        """Returns the values as an ``ndarray``."""
+        ...
     def is_i8(self) -> bool: ...
     def is_i16(self) -> bool: ...
     def is_f64(self) -> bool: ...

crates/python/qcs_sdk/qpu/__init__.pyi

@@ -29,6 +29,10 @@ class ReadoutValues:
 
     """
 
+    def __new__(cls, values: Union[List[int], List[float], List[complex]]):
+        """Construct a new ReadoutValues from a list of values."""
+        ...
+
     def inner(self) -> Union[List[int], List[float], List[complex]]:
         """Return the inner list of readout values."""
         ...

crates/python/src/execution_data.rs

@@ -1,23 +1,28 @@
 use std::time::Duration;
 
 use numpy::{Complex64, PyArray2};
-use pyo3::{exceptions::PyValueError, pymethods, types::PyDelta, Py, PyResult, Python};
+use pyo3::exceptions::PyKeyError;
+use pyo3::{
+    exceptions::PyValueError, pyclass, pymethods, types::PyDelta, Py, PyRef, PyRefMut, PyResult,
+    Python,
+};
+use pyo3::{IntoPy, PyAny};
 use qcs::{ExecutionData, RegisterMap, RegisterMatrix, ResultData};
-use qcs_api_client_grpc::models::controller::{readout_values::Values, ReadoutValues};
 use rigetti_pyo3::{
-    py_wrap_data_struct, py_wrap_error, py_wrap_type, py_wrap_union_enum, wrap_error, PyTryFrom,
-    PyWrapper, ToPython, ToPythonError,
+    impl_repr, py_wrap_data_struct, py_wrap_error, py_wrap_type, py_wrap_union_enum, wrap_error,
+    PyTryFrom, PyWrapper, ToPython, ToPythonError,
 };
 
+use crate::qpu::PyQpuResultData;
 use crate::qvm::PyQvmResultData;
-use crate::{grpc::models::controller::PyReadoutValuesValues, qpu::PyQpuResultData};
 
 py_wrap_union_enum! {
     PyResultData(ResultData) as "ResultData" {
         qpu: Qpu => PyQpuResultData,
         qvm: Qvm => PyQvmResultData
     }
 }
+impl_repr!(PyQpuResultData);
 
 wrap_error!(RustRegisterMatrixConversionError(
     qcs::RegisterMatrixConversionError
@@ -46,35 +51,53 @@ py_wrap_data_struct! {
         duration: Option<Duration> => Option<Py<PyDelta>>
     }
 }
+impl_repr!(PyExecutionData);
 
 #[pymethods]
 impl PyExecutionData {
     #[new]
-    fn __new__(py: Python<'_>, result_data: PyResultData, duration: Option<u64>) -> PyResult<Self> {
+    fn __new__(
+        py: Python<'_>,
+        result_data: PyResultData,
+        duration: Option<Py<PyDelta>>,
+    ) -> PyResult<Self> {
         Ok(Self(ExecutionData {
             result_data: ResultData::py_try_from(py, &result_data)?,
-            duration: duration.map(Duration::from_micros),
+            duration: match duration {
+                None => None,
+                Some(delta) => Some(
+                    delta
+                        .as_ref(py)
+                        .call_method0("total_seconds")
+                        .map(|result| result.extract::<f64>())?
+                        .map(Duration::from_secs_f64)?,
+                ),
+            },
         }))
     }
 }
 
-// From gRPC
-py_wrap_data_struct! {
-    PyReadoutValues(ReadoutValues) as "ReadoutValues" {
-        values: Option<Values> => Option<PyReadoutValuesValues>
-    }
-}
-
 py_wrap_type! {
+    #[pyo3(mapping)]
     PyRegisterMap(RegisterMap) as "RegisterMap";
 }
+impl_repr!(PyRegisterMap);
 
 py_wrap_type! {
     PyRegisterMatrix(RegisterMatrix) as "RegisterMatrix"
 }
+impl_repr!(PyRegisterMatrix);
 
 #[pymethods]
 impl PyRegisterMatrix {
+    fn to_ndarray(&self, py: Python<'_>) -> Py<PyAny> {
+        self.as_integer(py)
+            .map(|array| array.into_py(py))
+            .or(self.as_real(py).map(|array| array.into_py(py)))
+            .or(self.as_complex(py).map(|array| array.into_py(py)))
+            .expect("A RegisterMatrix can't be any other type.")
+    }
+
     #[staticmethod]
     fn from_integer(matrix: &PyArray2<i64>) -> PyRegisterMatrix {
         Self(RegisterMatrix::Integer(matrix.to_owned_array()))
@@ -153,9 +176,114 @@ impl PyRegisterMatrix {
 
 #[pymethods]
 impl PyRegisterMap {
-    pub fn get_register_matrix(&self, register_name: String) -> Option<PyRegisterMatrix> {
+    pub fn get_register_matrix(&self, register_name: &str) -> Option<PyRegisterMatrix> {
         self.as_inner()
-            .get_register_matrix(&register_name)
+            .get_register_matrix(register_name)
             .map(PyRegisterMatrix::from)
     }
+
+    pub fn __len__(&self) -> usize {
+        self.as_inner().0.len()
+    }
+
+    pub fn __contains__(&self, key: String) -> bool {
+        self.as_inner().0.contains_key(&key)
+    }
+
+    pub fn __getitem__(&self, item: &str) -> PyResult<PyRegisterMatrix> {
+        self.get_register_matrix(item)
+            .ok_or(PyKeyError::new_err(format!(
+                "Key {item} not found in RegisterMap"
+            )))
+    }
+
+    pub fn __iter__(&self, py: Python<'_>) -> PyResult<Py<PyRegisterMapKeysIter>> {
+        Py::new(
+            py,
+            PyRegisterMapKeysIter {
+                inner: self.as_inner().0.clone().into_iter(),
+            },
+        )
+    }
+
+    pub fn keys(&self, py: Python<'_>) -> PyResult<Py<PyRegisterMapKeysIter>> {
+        self.__iter__(py)
+    }
+
+    pub fn values(&self, py: Python<'_>) -> PyResult<Py<PyRegisterMapValuesIter>> {
+        Py::new(
+            py,
+            PyRegisterMapValuesIter {
+                inner: self.as_inner().0.clone().into_iter(),
+            },
+        )
+    }
+
+    pub fn items(&self, py: Python<'_>) -> PyResult<Py<PyRegisterMapItemsIter>> {
+        Py::new(
+            py,
+            PyRegisterMapItemsIter {
+                inner: self.as_inner().0.clone().into_iter(),
+            },
+        )
+    }
+
+    pub fn get(&self, key: &str, default: Option<PyRegisterMatrix>) -> Option<PyRegisterMatrix> {
+        self.__getitem__(key).ok().or(default)
+    }
+}
+
+#[pyclass]
+pub struct PyRegisterMapItemsIter {
+    inner: std::collections::hash_map::IntoIter<String, RegisterMatrix>,
+}
+
+#[pymethods]
+impl PyRegisterMapItemsIter {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<(String, PyRegisterMatrix)> {
+        slf.inner
+            .next()
+            .map(|(register, matrix)| (register, PyRegisterMatrix(matrix)))
+    }
+}
+
+// The keys and values iterators are built on the iterator of the full
+// `HashMap`, because the iterators returned by `keys()` and `values()`
+// return an iterator with a _reference_ to the underlying `HashMap`.
+// The reference would require these structs to specify a lifetime,
+// which is incompatible with `#[pyclass]`.
+#[pyclass]
+pub struct PyRegisterMapKeysIter {
+    inner: std::collections::hash_map::IntoIter<String, RegisterMatrix>,
+}
+
+#[pymethods]
+impl PyRegisterMapKeysIter {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<String> {
+        slf.inner.next().map(|(register, _)| register)
+    }
+}
+
+#[pyclass]
+pub struct PyRegisterMapValuesIter {
+    inner: std::collections::hash_map::IntoIter<String, RegisterMatrix>,
+}
+
+#[pymethods]
+impl PyRegisterMapValuesIter {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<PyRegisterMatrix> {
+        slf.inner.next().map(|(_, matrix)| PyRegisterMatrix(matrix))
+    }
 }

crates/python/src/grpc/models/controller.rs

@@ -1,11 +1,12 @@
+use numpy::PyArray;
 use pyo3::{
     prelude::*,
     types::{PyComplex, PyInt, PyList},
 };
 use qcs_api_client_grpc::models::controller::{
     readout_values::Values, Complex64, Complex64ReadoutValues, IntegerReadoutValues, ReadoutValues,
 };
-use rigetti_pyo3::{num_complex::Complex32 as NumComplex32, py_wrap_struct};
+use rigetti_pyo3::{num_complex::Complex32 as NumComplex32, py_wrap_struct, PyWrapper};
 use rigetti_pyo3::{py_wrap_data_struct, py_wrap_union_enum, PyTryFrom, ToPython};
 
 py_wrap_data_struct! {
@@ -14,6 +15,27 @@ py_wrap_data_struct! {
     }
 }
 
+#[pymethods]
+impl PyReadoutValues {
+    pub fn as_ndarray(&self, py: Python<'_>) -> PyResult<Option<PyObject>> {
+        match &self.as_inner().values {
+            None => Ok(None),
+            Some(Values::IntegerValues(ints)) => {
+                Ok(Some(PyArray::from_slice(py, &ints.values).to_object(py)))
+            }
+            Some(Values::ComplexValues(complex)) => Ok(Some(
+                PyArray::from_iter(py, {
+                    complex
+                        .values
+                        .iter()
+                        .map(|value| NumComplex32::new(value.real, value.imaginary))
+                })
+                .to_object(py),
+            )),
+        }
+    }
+}
+
 py_wrap_union_enum! {
     PyReadoutValuesValues(Values) as "ReadoutValuesValues" {
         integer_values: IntegerValues => PyIntegerReadoutValues,

crates/python/src/qpu/api.rs

@@ -158,6 +158,25 @@ impl From<readout_values::Values> for ExecutionResult {
     }
 }
 
+#[pymethods]
+impl ExecutionResult {
+    #[staticmethod]
+    fn from_register(register: PyRegister) -> Self {
+        match register.as_inner() {
+            Register::I32(values) => ExecutionResult {
+                shape: [values.len(), 1],
+                dtype: "integer".into(),
+                data: register,
+            },
+            Register::Complex32(values) => ExecutionResult {
+                shape: [values.len(), 1],
+                dtype: "complex".into(),
+                data: register,
+            },
+        }
+    }
+}
+
 #[pyclass]
 #[derive(Clone, Debug)]
 pub struct ExecutionResults {
@@ -169,6 +188,20 @@ pub struct ExecutionResults {
     pub execution_duration_microseconds: Option<u64>,
 }
 
+#[pymethods]
+impl ExecutionResults {
+    #[new]
+    fn new(
+        buffers: HashMap<String, ExecutionResult>,
+        execution_duration_microseconds: Option<u64>,
+    ) -> Self {
+        Self {
+            buffers,
+            execution_duration_microseconds,
+        }
+    }
+}
+
 impl From<ControllerJobExecutionResult> for ExecutionResults {
     fn from(value: ControllerJobExecutionResult) -> Self {
         let buffers = value

crates/python/src/qpu/result_data.rs

@@ -2,8 +2,8 @@ use std::collections::HashMap;
 
 use pyo3::{
     pymethods,
-    types::{PyComplex, PyFloat, PyInt},
-    Py, PyResult, Python,
+    types::{PyComplex, PyDict, PyFloat, PyInt},
+    Py, PyAny, PyResult, Python, ToPyObject,
 };
 use qcs::qpu::{QpuResultData, ReadoutValues};
 use rigetti_pyo3::{py_wrap_type, py_wrap_union_enum, PyTryFrom, PyWrapper, ToPython};
@@ -43,4 +43,27 @@ impl PyQpuResultData {
     fn readout_values(&self, py: Python<'_>) -> PyResult<HashMap<String, PyReadoutValues>> {
         self.as_inner().readout_values().to_python(py)
     }
+
+    fn asdict(&self, py: Python<'_>) -> PyResult<Py<PyDict>> {
+        let dict = PyDict::new(py);
+        dict.set_item("mappings", self.mappings(py)?)?;
+        dict.set_item(
+            "readout_values",
+            self.as_inner()
+                .readout_values()
+                .iter()
+                .map(|(register, values)| {
+                    (
+                        register.to_string(),
+                        match values {
+                            ReadoutValues::Integer(values) => values.to_object(py),
+                            ReadoutValues::Real(values) => values.to_object(py),
+                            ReadoutValues::Complex(values) => values.to_object(py),
+                        },
+                    )
+                })
+                .collect::<HashMap<String, Py<PyAny>>>(),
+        )?;
+        Ok(dict.into())
+    }
 }