feat(py): python bindings for BoolExpr with sugar

pyargus/src/lib.rs

@@ -1,9 +1,10 @@
-use argus_core::expr::{BoolExpr, NumExpr};
+use argus_core::expr::Ordering;
 use argus_core::prelude::*;
 use pyo3::prelude::*;
+use pyo3::pyclass::CompareOp;
 
 #[pyclass(name = "NumExpr", subclass)]
-#[derive(Clone)]
+#[derive(Clone, derive_more::From)]
 struct PyNumExpr(Box<NumExpr>);
 
 #[pymethods]
@@ -24,6 +25,10 @@ impl PyNumExpr {
         Python::with_gil(|py| Py::new(py, Mul::new(vec![self.clone(), other.clone()])))
     }
 
+    fn __sub__(&self, other: &Self) -> PyResult<Py<Sub>> {
+        Python::with_gil(|py| Py::new(py, Sub::new(self.clone(), other.clone())))
+    }
+
     fn __truediv__(&self, other: &Self) -> PyResult<Py<Div>> {
         Python::with_gil(|py| Py::new(py, Div::new(self.clone(), other.clone())))
     }
@@ -31,6 +36,17 @@ impl PyNumExpr {
     fn __abs__(&self) -> PyResult<Py<Abs>> {
         Python::with_gil(|py| Py::new(py, Abs::new(self.clone())))
     }
+
+    fn __richcmp__(&self, other: &Self, op: CompareOp) -> PyResult<Py<Cmp>> {
+        match op {
+            CompareOp::Lt => Cmp::less_than(self.clone(), other.clone()),
+            CompareOp::Le => Cmp::less_than_eq(self.clone(), other.clone()),
+            CompareOp::Eq => Cmp::equal(self.clone(), other.clone()),
+            CompareOp::Ne => Cmp::not_equal(self.clone(), other.clone()),
+            CompareOp::Gt => Cmp::greater_than(self.clone(), other.clone()),
+            CompareOp::Ge => Cmp::greater_than_eq(self.clone(), other.clone()),
+        }
+    }
 }
 
 #[pyclass(extends=PyNumExpr)]
@@ -40,7 +56,7 @@ struct ConstInt;
 impl ConstInt {
     #[new]
     fn new(val: i64) -> (Self, PyNumExpr) {
-        (Self, PyNumExpr(Box::new(NumExpr::IntLit(val))))
+        (Self, Box::new(NumExpr::IntLit(val)).into())
     }
 }
 
@@ -51,7 +67,7 @@ struct ConstUInt;
 impl ConstUInt {
     #[new]
     fn new(val: u64) -> (Self, PyNumExpr) {
-        (Self, PyNumExpr(Box::new(NumExpr::UIntLit(val))))
+        (Self, Box::new(NumExpr::UIntLit(val)).into())
     }
 }
 
@@ -62,7 +78,7 @@ struct ConstFloat;
 impl ConstFloat {
     #[new]
     fn new(val: f64) -> (Self, PyNumExpr) {
-        (Self, PyNumExpr(Box::new(NumExpr::FloatLit(val))))
+        (Self, Box::new(NumExpr::FloatLit(val)).into())
     }
 }
 
@@ -73,7 +89,7 @@ struct VarInt;
 impl VarInt {
     #[new]
     fn new(name: String) -> (Self, PyNumExpr) {
-        (Self, PyNumExpr(Box::new(NumExpr::IntVar { name })))
+        (Self, Box::new(NumExpr::IntVar { name }).into())
     }
 }
 
@@ -84,7 +100,7 @@ struct VarUInt;
 impl VarUInt {
     #[new]
     fn new(name: String) -> (Self, PyNumExpr) {
-        (Self, PyNumExpr(Box::new(NumExpr::UIntVar { name })))
+        (Self, Box::new(NumExpr::UIntVar { name }).into())
     }
 }
 
@@ -95,7 +111,7 @@ struct VarFloat;
 impl VarFloat {
     #[new]
     fn new(name: String) -> (Self, PyNumExpr) {
-        (Self, PyNumExpr(Box::new(NumExpr::FloatVar { name })))
+        (Self, Box::new(NumExpr::FloatVar { name }).into())
     }
 }
 
@@ -107,7 +123,7 @@ impl Negate {
     #[new]
     fn new(arg: PyNumExpr) -> (Self, PyNumExpr) {
         let arg = arg.0;
-        (Self, PyNumExpr(Box::new(NumExpr::Neg { arg })))
+        (Self, Box::new(NumExpr::Neg { arg }).into())
     }
 }
 
@@ -119,7 +135,7 @@ impl Add {
     #[new]
     fn new(args: Vec<PyNumExpr>) -> (Self, PyNumExpr) {
         let args: Vec<NumExpr> = args.into_iter().map(|arg| *arg.0).collect();
-        (Self, PyNumExpr(Box::new(NumExpr::Add { args })))
+        (Self, Box::new(NumExpr::Add { args }).into())
     }
 }
 
@@ -132,7 +148,7 @@ impl Sub {
     fn new(lhs: PyNumExpr, rhs: PyNumExpr) -> (Self, PyNumExpr) {
         let lhs = lhs.0;
         let rhs = rhs.0;
-        (Self, PyNumExpr(Box::new(NumExpr::Sub { lhs, rhs })))
+        (Self, Box::new(NumExpr::Sub { lhs, rhs }).into())
     }
 }
 
@@ -144,7 +160,7 @@ impl Mul {
     #[new]
     fn new(args: Vec<PyNumExpr>) -> (Self, PyNumExpr) {
         let args: Vec<NumExpr> = args.into_iter().map(|arg| *arg.0).collect();
-        (Self, PyNumExpr(Box::new(NumExpr::Mul { args })))
+        (Self, Box::new(NumExpr::Mul { args }).into())
     }
 }
 
@@ -157,23 +173,10 @@ impl Div {
     fn new(dividend: PyNumExpr, divisor: PyNumExpr) -> (Self, PyNumExpr) {
         let dividend = dividend.0;
         let divisor = divisor.0;
-        (Self, PyNumExpr(Box::new(NumExpr::Div { dividend, divisor })))
+        (Self, Box::new(NumExpr::Div { dividend, divisor }).into())
     }
 }
 
-// #[pyclass(extends=PyNumExpr)]
-// struct Pow;
-//
-// #[pymethods]
-// impl Pow {
-//     #[new]
-//     fn new(base: PyNumExpr, exponent: PyNumExpr) -> (Self, PyNumExpr) {
-//         let base = base.0;
-//         let exponent = exponent.0;
-//         (Self, PyNumExpr(Box::new(NumExpr:: { base, exponent })))
-//     }
-// }
-
 #[pyclass(extends=PyNumExpr)]
 struct Abs;
 
@@ -182,14 +185,141 @@ impl Abs {
     #[new]
     fn new(arg: PyNumExpr) -> (Self, PyNumExpr) {
         let arg = arg.0;
-        (Self, PyNumExpr(Box::new(NumExpr::Abs { arg })))
+        (Self, Box::new(NumExpr::Abs { arg }).into())
     }
 }
 
-#[pyclass(name = "BoolExpr")]
-struct PyBoolExpr(BoolExpr);
+#[pyclass(name = "BoolExpr", subclass)]
+#[derive(Clone, derive_more::From)]
+struct PyBoolExpr(Box<BoolExpr>);
+
+#[pymethods]
+impl PyBoolExpr {
+    fn __repr__(&self) -> String {
+        format!("{:?}", &self.0)
+    }
+
+    fn __invert__(&self) -> PyResult<Py<Not>> {
+        Python::with_gil(|py| Py::new(py, Not::new(self.clone())))
+    }
+
+    fn __or__(&self, other: &Self) -> PyResult<Py<Or>> {
+        Python::with_gil(|py| Py::new(py, Or::new(vec![self.clone(), other.clone()])))
+    }
+
+    fn __and__(&self, other: &Self) -> PyResult<Py<And>> {
+        Python::with_gil(|py| Py::new(py, And::new(vec![self.clone(), other.clone()])))
+    }
+}
+
+#[pyclass(extends=PyBoolExpr)]
+struct ConstBool;
+
+#[pymethods]
+impl ConstBool {
+    #[new]
+    fn new(val: bool) -> (Self, PyBoolExpr) {
+        (Self, Box::new(BoolExpr::BoolLit(val)).into())
+    }
+}
+
+#[pyclass(extends=PyBoolExpr)]
+struct VarBool;
+
+#[pymethods]
+impl VarBool {
+    #[new]
+    fn new(name: String) -> (Self, PyBoolExpr) {
+        (Self, Box::new(BoolExpr::BoolVar { name }).into())
+    }
+}
+
+#[pyclass(extends=PyBoolExpr)]
+struct Cmp;
+
+#[pyclass]
+#[derive(Copy, Clone, derive_more::From)]
+struct PyOrdering(Ordering);
+
+#[pymethods]
+impl Cmp {
+    #[new]
+    fn new(op: PyOrdering, lhs: PyNumExpr, rhs: PyNumExpr) -> (Self, PyBoolExpr) {
+        let op = op.0;
+        let lhs = lhs.0;
+        let rhs = rhs.0;
+        (Self, Box::new(BoolExpr::Cmp { op, lhs, rhs }).into())
+    }
+
+    #[staticmethod]
+    fn less_than(lhs: PyNumExpr, rhs: PyNumExpr) -> PyResult<Py<Self>> {
+        Python::with_gil(|py| Py::new(py, Cmp::new(PyOrdering(Ordering::less_than()), lhs, rhs)))
+    }
+
+    #[staticmethod]
+    fn less_than_eq(lhs: PyNumExpr, rhs: PyNumExpr) -> PyResult<Py<Self>> {
+        Python::with_gil(|py| Py::new(py, Cmp::new(PyOrdering(Ordering::less_than_eq()), lhs, rhs)))
+    }
+
+    #[staticmethod]
+    fn greater_than(lhs: PyNumExpr, rhs: PyNumExpr) -> PyResult<Py<Self>> {
+        Python::with_gil(|py| Py::new(py, Cmp::new(PyOrdering(Ordering::greater_than()), lhs, rhs)))
+    }
+
+    #[staticmethod]
+    fn greater_than_eq(lhs: PyNumExpr, rhs: PyNumExpr) -> PyResult<Py<Self>> {
+        Python::with_gil(|py| Py::new(py, Cmp::new(PyOrdering(Ordering::greater_than_eq()), lhs, rhs)))
+    }
+
+    #[staticmethod]
+    fn equal(lhs: PyNumExpr, rhs: PyNumExpr) -> PyResult<Py<Self>> {
+        Python::with_gil(|py| Py::new(py, Cmp::new(PyOrdering(Ordering::equal()), lhs, rhs)))
+    }
+
+    #[staticmethod]
+    fn not_equal(lhs: PyNumExpr, rhs: PyNumExpr) -> PyResult<Py<Self>> {
+        Python::with_gil(|py| Py::new(py, Cmp::new(PyOrdering(Ordering::not_equal()), lhs, rhs)))
+    }
+}
+
+#[pyclass(extends=PyBoolExpr)]
+struct Not;
+
+#[pymethods]
+impl Not {
+    #[new]
+    fn new(arg: PyBoolExpr) -> (Self, PyBoolExpr) {
+        let arg = arg.0;
+        (Self, PyBoolExpr(Box::new(BoolExpr::Not { arg })))
+    }
+}
+
+#[pyclass(extends=PyBoolExpr)]
+struct And;
+
+#[pymethods]
+impl And {
+    #[new]
+    fn new(args: Vec<PyBoolExpr>) -> (Self, PyBoolExpr) {
+        let args: Vec<BoolExpr> = args.into_iter().map(|arg| *arg.0).collect();
+        (Self, PyBoolExpr(Box::new(BoolExpr::And { args })))
+    }
+}
+
+#[pyclass(extends=PyBoolExpr)]
+struct Or;
+
+#[pymethods]
+impl Or {
+    #[new]
+    fn new(args: Vec<PyBoolExpr>) -> (Self, PyBoolExpr) {
+        let args: Vec<BoolExpr> = args.into_iter().map(|arg| *arg.0).collect();
+        (Self, PyBoolExpr(Box::new(BoolExpr::Or { args })))
+    }
+}
 
 #[pymodule]
+#[pyo3(name = "_argus")]
 fn pyargus(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyNumExpr>()?;
     m.add_class::<ConstInt>()?;
@@ -202,9 +332,14 @@ fn pyargus(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<Add>()?;
     m.add_class::<Mul>()?;
     m.add_class::<Div>()?;
-    // m.add_class::<Pow>()?;
 
     m.add_class::<PyBoolExpr>()?;
+    m.add_class::<ConstBool>()?;
+    m.add_class::<VarBool>()?;
+    m.add_class::<Cmp>()?;
+    m.add_class::<Not>()?;
+    m.add_class::<And>()?;
+    m.add_class::<Or>()?;
 
     Ok(())
 }