qcs_sdk/
register_data.rs

1use numpy::PyArray;
2use pyo3::{
3    pymethods,
4    types::{PyComplex, PyFloat, PyInt},
5    Py, PyErr, PyObject, PyResult, Python, ToPyObject,
6};
7use qcs::RegisterData;
8use rigetti_pyo3::{py_wrap_union_enum, PyWrapper};
9
10py_wrap_union_enum! {
11    PyRegisterData(RegisterData) as "RegisterData" {
12        i8: I8 => Vec<Vec<Py<PyInt>>>,
13        f64: F64 => Vec<Vec<Py<PyFloat>>>,
14        i16: I16 => Vec<Vec<Py<PyInt>>>,
15        complex32: Complex32 => Vec<Vec<Py<PyComplex>>>
16    }
17}
18
19#[pymethods]
20impl PyRegisterData {
21    pub fn as_ndarray(&self, py: Python<'_>) -> PyResult<PyObject> {
22        match self.as_inner() {
23            RegisterData::I8(matrix) => {
24                PyArray::from_vec2(py, matrix.as_slice()).map(|arr| arr.to_object(py))
25            }
26            RegisterData::F64(matrix) => {
27                PyArray::from_vec2(py, matrix.as_slice()).map(|arr| arr.to_object(py))
28            }
29            RegisterData::I16(matrix) => {
30                PyArray::from_vec2(py, matrix.as_slice()).map(|arr| arr.to_object(py))
31            }
32            RegisterData::Complex32(matrix) => {
33                PyArray::from_vec2(py, matrix.as_slice()).map(|arr| arr.to_object(py))
34            }
35        }
36        .map_err(PyErr::from)
37    }
38}