Merge pull request #170 from rigetti/163-frame-set-python-bindings

Python bindings for the FrameSet API

quil-py/quil/program/__init__.pyi

@@ -1,4 +1,4 @@
-from typing import Dict, Set, final, List, Optional
+from typing import Dict, FrozenSet, Set, final, List, Optional
 
 from quil.instructions import (
     AttributeValue,
@@ -155,6 +155,42 @@ class CalibrationSet:
 class FrameSet:
     @staticmethod
     def __new__(cls) -> "FrameSet": ...
+    def get(self, identifier: FrameIdentifier) -> Optional[Dict[str, AttributeValue]]:
+        """
+        Retrieve the attributes of a frame by its identifier
+        """
+        ...
+    def get_keys(self) -> List[FrameIdentifier]:
+        """
+        Return a list of all ``FrameIdentifier``s described by this ``FrameSet``
+        """
+        ...
+    def insert(
+        self, identifier: FrameIdentifier, attributes: Dict[str, AttributeValue]
+    ):
+        """
+        Insert a new ``FrameIdentifier``, overwriting any existing one.
+        """
+        ...
+    def merge(self, other: FrameSet):
+        """
+        Merge another ``FrameSet`` into this one, overwriting any existing keys.
+        """
+    def intersection(self, identifiers: FrozenSet[FrameIdentifier]) -> "FrameSet":
+        """
+        Return a new ``FrameSet`` which describes only the given ``FrameIdentifier``s
+        """
+        ...
+    def is_empty(self) -> bool:
+        """
+        Returns ``True`` if this ``FrameSet`` defines no frames.
+        """
+        ...
+    def to_instructions(self) -> List[Instruction]:
+        """
+        Return the Quil instructions which define the contained frames.
+        """
+        ...
     def get_all_frames(self) -> Dict[FrameIdentifier, Dict[str, AttributeValue]]: ...
 
 class MemoryRegion:

quil-py/src/program/frame.rs

@@ -1,19 +1,25 @@
-use std::collections::HashMap;
+use std::borrow::Borrow;
+use std::collections::{HashMap, HashSet};
+use std::convert::AsRef;
 
-use quil_rs::program::FrameSet;
+use quil_rs::{
+    instruction::{FrameAttributes, FrameIdentifier},
+    program::FrameSet,
+};
 use rigetti_pyo3::{
-    impl_repr, py_wrap_type,
+    impl_as_mut_for_wrapper, impl_repr, py_wrap_type,
     pyo3::{pyclass::CompareOp, pymethods, IntoPy, PyObject, PyResult, Python},
-    PyWrapper, ToPython,
+    PyTryFrom, PyWrapper, PyWrapperMut, ToPython,
 };
 
-use crate::instruction::{PyFrameAttributes, PyFrameIdentifier};
+use crate::instruction::{PyFrameAttributes, PyFrameIdentifier, PyInstruction};
 
 py_wrap_type! {
     #[derive(Debug, PartialEq, Eq)]
     PyFrameSet(FrameSet) as "FrameSet"
 }
 impl_repr!(PyFrameSet);
+impl_as_mut_for_wrapper!(PyFrameSet);
 
 #[pymethods]
 impl PyFrameSet {
@@ -22,6 +28,20 @@ impl PyFrameSet {
         Self(FrameSet::new())
     }
 
+    pub fn get(
+        &self,
+        py: Python<'_>,
+        identifier: PyFrameIdentifier,
+    ) -> PyResult<Option<PyFrameAttributes>> {
+        self.as_inner()
+            .get(&FrameIdentifier::py_try_from(py, &identifier)?)
+            .to_python(py)
+    }
+
+    pub fn get_keys(&self, py: Python<'_>) -> PyResult<Vec<PyFrameIdentifier>> {
+        self.as_inner().get_keys().to_python(py)
+    }
+
     pub fn get_all_frames(
         &self,
         py: Python<'_>,
@@ -32,6 +52,52 @@ impl PyFrameSet {
             .collect()
     }
 
+    pub fn insert(
+        &mut self,
+        py: Python<'_>,
+        identifier: PyFrameIdentifier,
+        attributes: PyFrameAttributes,
+    ) -> PyResult<()> {
+        self.as_inner_mut().insert(
+            FrameIdentifier::py_try_from(py, &identifier)?,
+            FrameAttributes::py_try_from(py, &attributes)?,
+        );
+        Ok(())
+    }
+
+    pub fn merge(&mut self, py: Python<'_>, other: PyFrameSet) -> PyResult<()> {
+        self.as_inner_mut()
+            .merge(FrameSet::py_try_from(py, &other)?);
+        Ok(())
+    }
+
+    pub fn intersection(
+        &self,
+        py: Python<'_>,
+        identifiers: HashSet<PyFrameIdentifier>,
+    ) -> PyResult<Self> {
+        Ok(Self(
+            self.as_inner().intersection(
+                &HashSet::<FrameIdentifier>::py_try_from(py, &identifiers)?
+                    .iter()
+                    .map(Borrow::borrow)
+                    .collect(),
+            ),
+        ))
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.as_inner().is_empty()
+    }
+
+    pub fn to_instructions(&self, py: Python<'_>) -> PyResult<Vec<PyInstruction>> {
+        self.as_inner().to_instructions().to_python(py)
+    }
+
+    pub fn __len__(&self) -> usize {
+        self.as_inner().len()
+    }
+
     pub fn __richcmp__(&self, py: Python<'_>, other: &Self, op: CompareOp) -> PyObject {
         match op {
             CompareOp::Eq => (self.as_inner() == other.as_inner()).into_py(py),

quil-rs/src/program/frame.rs

@@ -32,6 +32,11 @@ impl FrameSet {
         }
     }
 
+    /// Retrieve the attributes of a frame by its identifier.
+    pub fn get(&self, identifier: &FrameIdentifier) -> Option<&FrameAttributes> {
+        self.frames.get(identifier)
+    }
+
     /// Return a list of all frame IDs described by this FrameSet.
     pub fn get_keys(&self) -> Vec<&FrameIdentifier> {
         self.frames.keys().collect()
@@ -80,11 +85,6 @@ impl FrameSet {
         }
     }
 
-    /// Retrieve the attributes of a frame by its identifier.
-    pub fn get(&self, identifier: &FrameIdentifier) -> Option<&FrameAttributes> {
-        self.frames.get(identifier)
-    }
-
     /// Insert a new frame by ID, overwriting any existing one.
     pub fn insert(&mut self, identifier: FrameIdentifier, attributes: FrameAttributes) {
         self.frames.insert(identifier, attributes);