Makes conversions more ergonomic

Author: saulshanabrook

src/conversions.rs

@@ -0,0 +1,29 @@
+// Create wrappers around input types so that convert from pyobjects to them
+// and then from them to the egg_smol types
+//
+// Converts from Python classes we define in pure python so we can use dataclasses
+// to represent the input types
+use pyo3::prelude::*;
+
+pub struct WrappedVariant(egg_smol::ast::Variant);
+
+impl FromPyObject<'_> for WrappedVariant {
+    fn extract(obj: &'_ PyAny) -> PyResult<Self> {
+        Ok(WrappedVariant(egg_smol::ast::Variant {
+            name: obj.getattr("name")?.extract::<String>()?.into(),
+            cost: obj.getattr("cost")?.extract()?,
+            types: obj
+                .getattr("types")?
+                .extract::<Vec<String>>()?
+                .into_iter()
+                .map(|x| x.into())
+                .collect(),
+        }))
+    }
+}
+
+impl From<WrappedVariant> for egg_smol::ast::Variant {
+    fn from(other: WrappedVariant) -> Self {
+        other.0
+    }
+}

src/lib.rs

@@ -1,4 +1,6 @@
+mod conversions;
 mod error;
+use conversions::*;
 use error::*;
 use pyo3::prelude::*;
 
@@ -11,24 +13,6 @@ struct EGraph {
     egraph: egg_smol::EGraph,
 }
 
-// Convert a Python Variant object into a rust variable, by getting the attributes
-fn get_variant(obj: &PyAny) -> PyResult<egg_smol::ast::Variant> {
-    // TODO: Is there a way to do this more automatically?
-    Ok(egg_smol::ast::Variant {
-        name: obj
-            .getattr(pyo3::intern!(obj.py(), "name"))?
-            .extract::<String>()?
-            .into(),
-        cost: obj.getattr(pyo3::intern!(obj.py(), "cost"))?.extract()?,
-        types: obj
-            .getattr(pyo3::intern!(obj.py(), "types"))?
-            .extract::<Vec<String>>()?
-            .into_iter()
-            .map(|x| x.into())
-            .collect(),
-    })
-}
-
 #[pymethods]
 impl EGraph {
     #[new]
@@ -53,12 +37,8 @@ impl EGraph {
     /// --
     ///
     /// Declare a new datatype constructor.
-    fn declare_constructor(
-        &mut self,
-        #[pyo3(from_py_with = "get_variant")] variant: egg_smol::ast::Variant,
-        sort: &str,
-    ) -> EggResult<()> {
-        self.egraph.declare_constructor(variant, sort)?;
+    fn declare_constructor(&mut self, variant: WrappedVariant, sort: &str) -> EggResult<()> {
+        self.egraph.declare_constructor(variant.into(), sort)?;
         Ok({})
     }