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/.stubtest-allowlist

@@ -1,3 +1,3 @@
-qcs_sdk.qcs_sdk
 qcs_sdk._gather_diagnostics
 qcs_sdk.diagnostics
+qcs_sdk.qcs_sdk

crates/python/qcs_sdk/__init__.pyi

@@ -1,6 +1,6 @@
 import datetime
 from enum import Enum
-from typing import Dict, List, Sequence, Optional, Union, final
+from typing import Dict, List, Sequence, Optional, Union, final, Iterable, Tuple
 
 import numpy as np
 from numpy.typing import NDArray
@@ -205,6 +205,11 @@ class RegisterMatrix:
     def from_real(inner: NDArray[np.float64]) -> "RegisterMatrix": ...
     @staticmethod
     def from_complex(inner: NDArray[np.complex128]) -> "RegisterMatrix": ...
+    def to_ndarray(self) -> Union[NDArray[np.complex128], NDArray[np.int64], NDArray[np.float64]]:
+        """
+        Get the RegisterMatrix as numpy ``ndarray``.
+        """
+        ...
 
 @final
 class RegisterMap:
@@ -213,6 +218,14 @@ class RegisterMap:
     def get_register_matrix(self, register_name: str) -> Optional[RegisterMatrix]:
         """Get the ``RegisterMatrix`` for the given register. Returns `None` if the register doesn't exist."""
         ...
+    def get(self, default: Optional[RegisterMatrix] = None) -> Optional[RegisterMatrix]: ...
+    def items(self) -> Iterable[Tuple[str, RegisterMatrix]]: ...
+    def keys(self) -> Iterable[str]: ...
+    def values(self) -> Iterable[RegisterMatrix]: ...
+    def __iter__(self) -> Iterable[str]: ...
+    def __getitem__(self, item: str) -> RegisterMatrix: ...
+    def __contains__(self, key: str) -> bool: ...
+    def __len__(self) -> int: ...
 
 @final
 class ResultData:
@@ -261,6 +274,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 +317,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 +351,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

@@ -1,4 +1,4 @@
-from typing import Dict, List, Mapping, Sequence, Optional, Union, final
+from typing import Dict, List, Mapping, Sequence, Optional, Union, final, Any
 
 from qcs_sdk.client import QCSClient
 
@@ -29,6 +29,9 @@ 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."""
         ...
@@ -52,11 +55,14 @@ class ReadoutValues:
 class QPUResultData:
     """
     Encapsulates data returned from the QPU after executing a job.
+
+    ``QPUResultData`` contains "mappings", which map declared memory regions
+    in a program (ie. "ro[0]") to that regions readout key in "readout_values".
+    "readout_values" maps those readout keys to the values emitted for that region
+    across all shots.
     """
 
-    def __new__(
-        cls, mappings: Mapping[str, str], readout_values: Mapping[str, ReadoutValues]
-    ): ...
+    def __new__(cls, mappings: Mapping[str, str], readout_values: Mapping[str, ReadoutValues]): ...
     @property
     def mappings(self) -> Dict[str, str]:
         """
@@ -69,6 +75,27 @@ class QPUResultData:
         Get the mappings of a readout values identifier (ie. "q0") to a set of ``ReadoutValues``
         """
         ...
+    def to_raw_readout_data(
+        self,
+    ) -> RawQPUReadoutData:
+        """
+        Get a copy of this result data flattened into a ``RawQPUReadoutData``
+        """
+        ...
+
+class RawQPUReadoutData:
+    @property
+    def mapping(self) -> Dict[str, str]:
+        """
+        Get the mappings of a memory region (ie. "ro[0]") to it's key name in readout_values
+        """
+        ...
+    @property
+    def readout_values(self) -> Dict[str, Union[List[int], List[float], List[complex]]]:
+        """
+        Get the mappings of a readout values identifier (ie. "q0") to a list of those readout values
+        """
+        ...
 
 class ListQuantumProcessorsError(RuntimeError):
     """A request to list available Quantum Processors failed."""

crates/python/qcs_sdk/qpu/api.pyi

@@ -48,6 +48,9 @@ class Register:
 class ExecutionResult:
     """Execution readout data from a particular memory location."""
 
+    @staticmethod
+    def from_register(register: Register) -> "ExecutionResult":
+        """Build an `ExecutionResult` from a `Register`."""
     @property
     def shape(self) -> List[int]:
         """The shape of the result data."""

crates/python/qcs_sdk/qvm/__init__.pyi

@@ -1,4 +1,4 @@
-from typing import final, Mapping, Optional, Sequence, Tuple, Union
+from typing import final, Mapping, Optional, Sequence, Tuple, Union, Dict, List
 
 from qcs_sdk import RegisterData, QCSClient
 
@@ -24,6 +24,23 @@ class QVMResultData:
         Get the mapping of register names (ie. "ro") to a ``RegisterData`` matrix containing the register values.
         """
         ...
+    @property
+    def to_raw_readout_data(
+        self,
+    ) -> RawQVMReadoutData:
+        """
+        Get a copy of this result data flattened into a ``RawQVMReadoutData``
+        """
+        ...
+
+@final
+class RawQVMReadoutData:
+    @property
+    def memory(self) -> Dict[str, Union[List[List[int]], List[List[float]]]]:
+        """
+        The mapping of register names (ie. "ro") to a 2-d list containing the
+        values for that register.
+        """
 
 @final
 class QVMOptions:

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))
+    }
 }