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}