1use std::borrow::Borrow;
2use std::collections::{HashMap, HashSet};
3use std::convert::AsRef;
4
5use quil_rs::{
6 instruction::{FrameAttributes, FrameIdentifier},
7 program::FrameSet,
8};
9use rigetti_pyo3::{
10 impl_as_mut_for_wrapper, impl_repr, py_wrap_type,
11 pyo3::{pymethods, PyResult, Python},
12 PyTryFrom, PyWrapper, PyWrapperMut, ToPython,
13};
14
15use crate::{
16 impl_eq,
17 instruction::{PyFrameAttributes, PyFrameIdentifier, PyInstruction},
18};
19
20py_wrap_type! {
21 #[derive(Debug, PartialEq, Eq)]
22 PyFrameSet(FrameSet) as "FrameSet"
23}
24impl_repr!(PyFrameSet);
25impl_as_mut_for_wrapper!(PyFrameSet);
26impl_eq!(PyFrameSet);
27
28#[pymethods]
29impl PyFrameSet {
30 #[new]
31 pub fn new() -> Self {
32 Self(FrameSet::new())
33 }
34
35 pub fn get(
36 &self,
37 py: Python<'_>,
38 identifier: PyFrameIdentifier,
39 ) -> PyResult<Option<PyFrameAttributes>> {
40 self.as_inner()
41 .get(&FrameIdentifier::py_try_from(py, &identifier)?)
42 .to_python(py)
43 }
44
45 pub fn get_keys(&self, py: Python<'_>) -> PyResult<Vec<PyFrameIdentifier>> {
46 self.as_inner().get_keys().to_python(py)
47 }
48
49 pub fn get_all_frames(
50 &self,
51 py: Python<'_>,
52 ) -> PyResult<HashMap<PyFrameIdentifier, PyFrameAttributes>> {
53 self.as_inner()
54 .iter()
55 .map(|(ident, attribs)| Ok((ident.to_python(py)?, attribs.to_python(py)?)))
56 .collect()
57 }
58
59 pub fn insert(
60 &mut self,
61 py: Python<'_>,
62 identifier: PyFrameIdentifier,
63 attributes: PyFrameAttributes,
64 ) -> PyResult<()> {
65 self.as_inner_mut().insert(
66 FrameIdentifier::py_try_from(py, &identifier)?,
67 FrameAttributes::py_try_from(py, &attributes)?,
68 );
69 Ok(())
70 }
71
72 pub fn merge(&mut self, py: Python<'_>, other: PyFrameSet) -> PyResult<()> {
73 self.as_inner_mut()
74 .merge(FrameSet::py_try_from(py, &other)?);
75 Ok(())
76 }
77
78 pub fn intersection(
79 &self,
80 py: Python<'_>,
81 identifiers: HashSet<PyFrameIdentifier>,
82 ) -> PyResult<Self> {
83 Ok(Self(
84 self.as_inner().intersection(
85 &HashSet::<FrameIdentifier>::py_try_from(py, &identifiers)?
86 .iter()
87 .map(Borrow::borrow)
88 .collect(),
89 ),
90 ))
91 }
92
93 pub fn is_empty(&self) -> bool {
94 self.as_inner().is_empty()
95 }
96
97 pub fn to_instructions(&self, py: Python<'_>) -> PyResult<Vec<PyInstruction>> {
98 self.as_inner().to_instructions().to_python(py)
99 }
100
101 pub fn __len__(&self) -> usize {
102 self.as_inner().len()
103 }
104}
105
106impl Default for PyFrameSet {
107 fn default() -> Self {
108 Self::new()
109 }
110}