Support user defined functions (#2)

Author: hardbyte

Cargo.toml

@@ -9,8 +9,8 @@ name = "cel"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = { version = "0.22.5", features = ["chrono", "gil-refs"]}
-cel-interpreter = "0.8.1"
+pyo3 = { version = "0.22.6", features = ["chrono", "gil-refs", "py-clone"]}
+cel-interpreter = {  version = "0.9.0", features = ["chrono", "json", "regex"] }
 log = "0.4.22"
 pyo3-log = "0.11.0"
 chrono = { version = "0.4.38", features = ["serde"] }

README.md

@@ -37,22 +37,23 @@ evaluate(
 )
 True
 ```
-## Future work
 
+### Custom Python Functions
 
-### Command line interface
+This Python library supports user defined Python functions
+in the context:
 
-The package (plans to) provides a command line interface for evaluating CEL expressions:
+```python
+from cel import evaluate
 
-```bash
-$ python -m cel '1 + 2'
-3
-```
+def is_adult(age):
+    return age > 21
 
-### Separate compilation and Execution steps
-### Custom Python Functions
+evaluate("is_adult(age)", {'is_adult': is_adult, 'age': 18})
+# False
+```
 
-Ability to add Python functions to the Context object:
+You can also explicitly create a Context object:
 
 ```python
 from cel import evaluate, Context
@@ -62,5 +63,24 @@ def is_adult(age):
 
 context = Context()
 context.add_function("is_adult", is_adult)
-print(evaluate("is_adult(age)", {"age": 18}, context))  # False
+context.update({"age": 18})
+
+evaluate("is_adult(age)", context)
+# False
 ```
+
+
+## Future work
+
+
+### Command line interface
+
+The package (plans to) provides a command line interface for evaluating CEL expressions:
+
+```bash
+$ python -m cel '1 + 2'
+3
+```
+
+### Separate compilation and Execution steps
+

src/context.rs

@@ -0,0 +1,85 @@
+use cel_interpreter::objects::TryIntoValue;
+use cel_interpreter::Value;
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyTuple};
+use std::collections::HashMap;
+use pyo3::exceptions::PyValueError;
+
+#[pyo3::pyclass]
+pub struct Context {
+    pub variables: HashMap<String, Value>,
+    pub functions: HashMap<String, Py<PyAny>>,
+}
+
+#[pyo3::pymethods]
+impl Context {
+
+    #[new]
+    pub fn new(variables: Option<&PyDict>, functions: Option<&PyDict>) -> PyResult<Self> {
+        let mut context = Context {
+            variables: HashMap::new(),
+            functions: HashMap::new(),
+        };
+
+        if let Some(variables) = variables {
+            //context.variables.extend(variables.clone());
+            for (k, v) in variables {
+                let key = k.extract::<String>().map_err(|_| {
+                    PyValueError::new_err("Keys must be strings")
+                });
+                key.map(|key| context.add_variable(key, v))??;
+
+            }
+        };
+
+        if let Some(functions) = functions {
+            context.update(functions)?;
+        };
+
+
+
+        Ok(context)
+    }
+
+
+    fn add_function(&mut self, name: String, function: Py<PyAny>) {
+        self.functions.insert(name, function);
+    }
+
+    pub fn add_variable(&mut self, name: String, value: &PyAny) -> PyResult<()> {
+        let value = crate::RustyPyType(value).try_into_value().map_err(|e| {
+            pyo3::exceptions::PyValueError::new_err(format!(
+                "Failed to convert variable '{}': {}",
+                name, e
+            ))
+        })?;
+        self.variables.insert(name, value);
+        Ok(())
+    }
+
+    pub fn update(&mut self, variables: &PyDict) -> PyResult<()> {
+
+        for (key, value) in variables {
+            // Attempt to extract the key as a String
+            let key = key.extract::<String>().map_err(|_| {
+                PyValueError::new_err("Keys must be strings")
+            })?;
+
+            if value.is_callable() {
+                // Value is a function, add it to the functions hashmap
+                let py_function = value.to_object(value.py());
+                self.functions.insert(key, py_function);
+            } else {
+                // Value is a variable, add it to the variables hashmap
+                let value = crate::RustyPyType(value)
+                    .try_into_value()
+                    .map_err(|e| PyValueError::new_err(e.to_string()))?;
+
+                self.variables.insert(key, value);
+            }
+
+        }
+
+        Ok(())
+    }
+}

src/lib.rs

@@ -1,17 +1,20 @@
+mod context;
+
 use cel_interpreter::objects::{Key, TryIntoValue};
-use cel_interpreter::{Context, Program, Value};
+use cel_interpreter::{ExecutionError, Program, Value};
 use log::{debug, info, warn};
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 
 use chrono::{DateTime, Duration as ChronoDuration, Offset, TimeZone, Utc};
-use pyo3::types::PyDelta;
 use pyo3::types::{PyBytes, PyDateTime, PyDict, PyList, PyTuple};
+use pyo3::types::{PyDelta, PyFunction};
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use std::collections::HashMap;
 use std::error::Error;
 use std::fmt;
+use std::ops::Deref;
 use std::sync::Arc;
 
 #[derive(Debug)]
@@ -183,79 +186,109 @@ impl TryIntoValue for RustyPyType<'_> {
 fn evaluate(src: String, evaluation_context: Option<&PyAny>) -> PyResult<RustyCelType> {
     debug!("Evaluating CEL expression: {}", src);
 
-    let context: Option<&PyDict> = evaluation_context.map(|context| {
-        context
-            .downcast::<PyDict>()
-            .expect("Failed to downcast PyDict")
-    });
+    let program = Program::compile(&src)
+        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Failed to compile expression '{}': {}", src, e)))?;
 
-    let program = Program::compile(src.as_str());
+    debug!("Compiled program: {:?}", program);
 
-    // Handle the result of the compilation
-    match program {
-        Err(compile_error) => {
-            debug!("An error occurred during compilation");
-            debug!("compile_error: {:?}", compile_error);
-            // compile_error
-            //     .into_iter()
-            //     .for_each(|e| println!("Parse error: {:?}", e));
-            return Err(PyValueError::new_err("Parse Error"));
+    debug!("Preparing context");
+    let mut environment = cel_interpreter::Context::default();
+    let mut ctx = context::Context::new(None, None)?;
+
+    // Custom Rust functions can also be added to the environment...
+    //environment.add_function("add", |a: i64, b: i64| a + b);
+
+    // Process the evaluation context if provided
+    if let Some(evaluation_context) = evaluation_context {
+        // Attempt to extract directly as a Context object
+        if let Ok(py_context_ref) = evaluation_context.extract::<PyRef<context::Context>>() {
+            // Clone variables and functions into our local Context
+            ctx.variables = py_context_ref.variables.clone();
+            ctx.functions = py_context_ref.functions.clone();
+
+        } else if let Ok(py_dict) = evaluation_context.extract::<&PyDict>() {
+            // User passed in a dict - let's process variables and functions from the dict
+            ctx.update(&py_dict)?;
+        } else {
+            return Err(PyValueError::new_err("evaluation_context must be a Context object or a dict"))
+        };
+
+
+        // Add any variables from the passed in Python context
+        for (name, value) in &ctx.variables {
+            environment
+                .add_variable(name.clone(), value.clone())
+                .map_err(|e| PyValueError::new_err(format!("Failed to add variable '{}': {}", name, e)))?;
         }
-        Ok(program) => {
-            let mut environment = Context::default();
-
-            // Custom functions can be added to the environment
-            //environment.add_function("add", |a: i64, b: i64| a + b);
-
-            // Add any variables from the passed in Dict context
-            if let Some(context) = context {
-                for (key, value) in context {
-                    debug!("Adding context {:?}", key);
-                    let key = key.extract::<String>().unwrap();
-                    // Each value is of type PyAny, we need to try to extract into a Value
-                    // and then add it to the CEL context
-
-                    let wrapped_value = RustyPyType(value);
-                    match wrapped_value.try_into_value() {
-                        Ok(value) => {
-                            debug!("Converted value: {:?}", value);
-                            environment
-                                .add_variable(key, value)
-                                .expect("Failed to add variable to context");
-                        }
-                        Err(error) => {
-                            debug!("An error occurred during context conversion");
-                            warn!("Conversion error: {:?}", error);
-                            warn!("Key: {:?}", key);
 
-                            return Err(PyValueError::new_err(error.to_string()));
+        // Add functions
+        let collected_functions: Vec<(String, Py<PyAny>)> = Python::with_gil(|py| {
+            ctx.functions
+                .iter()
+                .map(|(name, py_function)| (name.clone(), py_function.clone_ref(py)))
+                .collect()
+        });
+
+        for (name, py_function) in collected_functions.into_iter() {
+            environment.add_function(
+                &name.clone(),
+                move |ftx: &cel_interpreter::FunctionContext| -> cel_interpreter::ResolveResult {
+                    Python::with_gil(|py| {
+                        // Convert arguments from Expression in ftx.args to PyObjects
+                        let mut py_args = Vec::new();
+                        for arg_expr in &ftx.args {
+                            let arg_value = ftx.ptx.resolve(arg_expr)?;
+                            let py_arg = RustyCelType(arg_value).into_py(py);
+                            py_args.push(py_arg);
                         }
-                    }
-                }
-            }
+                        let py_args = PyTuple::new_bound(py, py_args);
 
-            let result = program.execute(&environment);
-            match result {
-                Err(error) => {
-                    warn!("An error occurred during execution");
-                    warn!("Execution error: {:?}", error);
-                    // errors
-                    //     .into_iter()
-                    //     .for_each(|e| println!("Execution error: {:?}", e));
-                    Err(PyValueError::new_err("Execution Error"))
-                }
+                        // Call the Python function
+                        let py_result = py_function.call1(py, py_args)
+                            .map_err(|e|  ExecutionError::FunctionError {
+                                function: name.clone(),
+                                message: e.to_string(),
+                            })?;
+                        // Convert the PyObject to &PyAny
+                        let py_result_ref = py_result.as_ref(py);
 
-                Ok(value) => return Ok(RustyCelType(value)),
-            }
+                        // Convert the result back to Value
+                        let value = RustyPyType(py_result_ref).try_into_value().map_err(|e| {
+                            ExecutionError::FunctionError {
+                                function: name.clone(),
+                                message: format!("Error calling function '{}': {}", name, e),
+                            }
+                        })?;
+                        Ok(value)
+                    })
+                },
+            );
         }
     }
+
+
+    let result = program.execute(&environment);
+    match result {
+        Err(error) => {
+            warn!("An error occurred during execution");
+            warn!("Execution error: {:?}", error);
+            // errors
+            //     .into_iter()
+            //     .for_each(|e| println!("Execution error: {:?}", e));
+            Err(PyValueError::new_err("Execution Error"))
+        }
+
+        Ok(value) => return Ok(RustyCelType(value)),
+    }
 }
 
 /// A Python module implemented in Rust.
 #[pymodule]
-fn cel(py: Python<'_>, m: &PyModule) -> PyResult<()> {
+fn cel<'py>(py: Python<'py>, m: &PyModule) -> PyResult<()> {
     pyo3_log::init();
 
     m.add_function(wrap_pyfunction!(evaluate, m)?)?;
+
+    m.add_class::<context::Context>()?;
     Ok(())
 }

tests/test_basics.py

@@ -9,6 +9,15 @@ def test_invalid_expression_raises_parse_value_error():
         result = cel.evaluate("1 +")
 
 
+def test_readme_example():
+    assert cel.evaluate(
+        'resource.name.startsWith("/groups/" + claim.group)',
+        {
+            "resource": {"name": "/groups/hardbyte"},
+            "claim": {"group": "hardbyte"}
+        }
+    )
+
 def test_hello_world():
     assert cel.evaluate("'Hello ' + name", {'name': "World"}) == "Hello World"