Files
rhai-python/src/lib.rs

331 lines
14 KiB
Rust

use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
use pyo3::IntoPyObjectExt;
use rhai::{Array, Dynamic, Engine, Map, Scope};
/// Convert Rhai Dynamic to Python object
fn dynamic_to_python(py: Python<'_>, value: Dynamic) -> PyResult<PyObject> {
match value.type_name() {
"i64" => Ok(value.as_int().unwrap().into_py_any(py).unwrap()),
"f64" => Ok(value.as_float().unwrap().into_py_any(py).unwrap()),
"bool" => Ok(value.as_bool().unwrap().into_py_any(py).unwrap()),
"string" => Ok(value.into_string().unwrap().into_py_any(py).unwrap()),
"char" => Ok(value
.as_char()
.unwrap()
.to_string()
.into_py_any(py)
.unwrap()),
"()" => Ok(py.None()),
"array" => {
let array = value.cast::<Array>();
let py_list = PyList::empty(py);
for item in array {
py_list.append(dynamic_to_python(py, item)?)?;
}
Ok(py_list.into_py_any(py)?.into())
}
"map" => {
let map = value.cast::<Map>();
let py_dict = PyDict::new(py);
for (key, value) in map {
let py_key = key.to_string().into_py_any(py).unwrap();
let py_value = dynamic_to_python(py, value)?;
py_dict.set_item(py_key, py_value)?;
}
Ok(py_dict.into_py_any(py)?.into())
}
_ => Err(PyRuntimeError::new_err(format!(
"Unsupported Rhai type: {}",
value.type_name()
))),
}
}
/// Convert Python object to Rhai Dynamic
fn python_to_dynamic(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<Dynamic> {
if obj.is_none() {
Ok(Dynamic::UNIT)
} else if let Ok(val) = obj.extract::<bool>() {
Ok(Dynamic::from(val))
} else if let Ok(val) = obj.extract::<i64>() {
Ok(Dynamic::from(val))
} else if let Ok(val) = obj.extract::<f64>() {
Ok(Dynamic::from(val))
} else if let Ok(val) = obj.extract::<String>() {
Ok(Dynamic::from(val))
} else if let Ok(py_list) = obj.downcast::<PyList>() {
let mut array = Array::new();
for item in py_list {
array.push(python_to_dynamic(py, &item)?);
}
Ok(Dynamic::from(array))
} else if let Ok(py_dict) = obj.downcast::<PyDict>() {
let mut map = Map::new();
for (key, value) in py_dict {
let key_str = key.extract::<String>()?;
let rhai_value = python_to_dynamic(py, &value)?;
map.insert(key_str.into(), rhai_value);
}
Ok(Dynamic::from(map))
} else {
Err(PyRuntimeError::new_err(format!(
"Unsupported Python type: {}",
obj.get_type().name()?
)))
}
}
#[pyclass]
struct RhaiEngine {
engine: Engine,
scope: Scope<'static>,
py_functions: std::collections::HashMap<String, PyObject>,
}
#[pymethods]
impl RhaiEngine {
#[new]
fn new() -> Self {
Self {
engine: Engine::new(),
scope: Scope::new(),
py_functions: std::collections::HashMap::new(),
}
}
/// Evaluate a Rhai script and return the result
fn eval(&mut self, py: Python<'_>, script: &str) -> PyResult<PyObject> {
match self
.engine
.eval_with_scope::<Dynamic>(&mut self.scope, script)
{
Ok(result) => dynamic_to_python(py, result),
Err(err) => Err(PyRuntimeError::new_err(format!("Rhai error: {}", err))),
}
}
/// Set a variable in the engine scope
fn set_var(&mut self, py: Python<'_>, name: &str, value: &Bound<'_, PyAny>) -> PyResult<()> {
let dynamic_value = python_to_dynamic(py, value)?;
self.scope.set_value(name, dynamic_value);
Ok(())
}
/// Get a variable from the engine scope
fn get_var(&self, py: Python<'_>, name: &str) -> PyResult<PyObject> {
match self.scope.get_value::<Dynamic>(name) {
Some(value) => dynamic_to_python(py, value),
None => Err(PyRuntimeError::new_err(format!(
"Variable '{}' not found",
name
))),
}
}
/// Check if a variable exists in the scope
fn has_var(&self, name: &str) -> bool {
self.scope.contains(name)
}
/// Clear all variables from the scope
fn clear_scope(&mut self) {
self.scope.clear();
}
/// Compile a script and return whether it's valid
fn compile(&self, script: &str) -> PyResult<bool> {
match self.engine.compile(script) {
Ok(_) => Ok(true),
Err(err) => Err(PyRuntimeError::new_err(format!(
"Compilation error: {}",
err
))),
}
}
/// Get the list of all variable names in the scope
fn list_vars(&self) -> Vec<String> {
self.scope
.iter()
.map(|(name, _, _)| name.to_string())
.collect()
}
/// Register a Python function to be callable from Rhai using signature inspection
fn register_fn(&mut self, py: Python<'_>, name: &str, func: PyObject) -> PyResult<()> {
// Get the function signature using inspect.signature()
let inspect_module = py.import("inspect")?;
let signature_func = inspect_module.getattr("signature")?;
let sig = signature_func.call1((&func,))?;
let parameters = sig.getattr("parameters")?;
let param_count = parameters.len()?;
// Store the function
self.py_functions
.insert(name.to_string(), func.clone_ref(py));
let func_name = name.to_string();
// Register the function with the appropriate arity
match param_count {
0 => {
self.engine.register_fn(name, {
let func = func.clone_ref(py);
let name = func_name.clone();
move || -> Result<Dynamic, Box<rhai::EvalAltResult>> {
Python::with_gil(|py| {
let result = func
.call0(py)
.map_err(|e| format!("Python function '{}' error: {}", name, e))?;
python_to_dynamic(py, &result.bind(py))
.map_err(|e| format!("Type conversion error: {}", e).into())
})
}
});
}
1 => {
self.engine.register_fn(name, {
let func = func.clone_ref(py);
let name = func_name.clone();
move |arg1: Dynamic| -> Result<Dynamic, Box<rhai::EvalAltResult>> {
Python::with_gil(|py| {
let py_arg1 = dynamic_to_python(py, arg1)
.map_err(|e| format!("Type conversion error: {}", e))?;
let result = func
.call1(py, (py_arg1,))
.map_err(|e| format!("Python function '{}' error: {}", name, e))?;
python_to_dynamic(py, &result.bind(py))
.map_err(|e| format!("Type conversion error: {}", e).into())
})
}
});
}
2 => {
self.engine.register_fn(name, {
let func = func.clone_ref(py);
let name = func_name.clone();
move |arg1: Dynamic,
arg2: Dynamic|
-> Result<Dynamic, Box<rhai::EvalAltResult>> {
Python::with_gil(|py| {
let py_arg1 = dynamic_to_python(py, arg1)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg2 = dynamic_to_python(py, arg2)
.map_err(|e| format!("Type conversion error: {}", e))?;
let result = func
.call1(py, (py_arg1, py_arg2))
.map_err(|e| format!("Python function '{}' error: {}", name, e))?;
python_to_dynamic(py, &result.bind(py))
.map_err(|e| format!("Type conversion error: {}", e).into())
})
}
});
}
3 => {
self.engine.register_fn(name, {
let func = func.clone_ref(py);
let name = func_name.clone();
move |arg1: Dynamic,
arg2: Dynamic,
arg3: Dynamic|
-> Result<Dynamic, Box<rhai::EvalAltResult>> {
Python::with_gil(|py| {
let py_arg1 = dynamic_to_python(py, arg1)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg2 = dynamic_to_python(py, arg2)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg3 = dynamic_to_python(py, arg3)
.map_err(|e| format!("Type conversion error: {}", e))?;
let result = func
.call1(py, (py_arg1, py_arg2, py_arg3))
.map_err(|e| format!("Python function '{}' error: {}", name, e))?;
python_to_dynamic(py, &result.bind(py))
.map_err(|e| format!("Type conversion error: {}", e).into())
})
}
});
}
4 => {
self.engine.register_fn(name, {
let func = func.clone_ref(py);
let name = func_name.clone();
move |arg1: Dynamic,
arg2: Dynamic,
arg3: Dynamic,
arg4: Dynamic|
-> Result<Dynamic, Box<rhai::EvalAltResult>> {
Python::with_gil(|py| {
let py_arg1 = dynamic_to_python(py, arg1)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg2 = dynamic_to_python(py, arg2)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg3 = dynamic_to_python(py, arg3)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg4 = dynamic_to_python(py, arg4)
.map_err(|e| format!("Type conversion error: {}", e))?;
let result = func
.call1(py, (py_arg1, py_arg2, py_arg3, py_arg4))
.map_err(|e| format!("Python function '{}' error: {}", name, e))?;
python_to_dynamic(py, &result.bind(py))
.map_err(|e| format!("Type conversion error: {}", e).into())
})
}
});
}
5 => {
self.engine.register_fn(name, {
let func = func.clone_ref(py);
let name = func_name.clone();
move |arg1: Dynamic,
arg2: Dynamic,
arg3: Dynamic,
arg4: Dynamic,
arg5: Dynamic|
-> Result<Dynamic, Box<rhai::EvalAltResult>> {
Python::with_gil(|py| {
let py_arg1 = dynamic_to_python(py, arg1)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg2 = dynamic_to_python(py, arg2)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg3 = dynamic_to_python(py, arg3)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg4 = dynamic_to_python(py, arg4)
.map_err(|e| format!("Type conversion error: {}", e))?;
let py_arg5 = dynamic_to_python(py, arg5)
.map_err(|e| format!("Type conversion error: {}", e))?;
let result = func
.call1(py, (py_arg1, py_arg2, py_arg3, py_arg4, py_arg5))
.map_err(|e| format!("Python function '{}' error: {}", name, e))?;
python_to_dynamic(py, &result.bind(py))
.map_err(|e| format!("Type conversion error: {}", e).into())
})
}
});
}
_ => {
return Err(PyRuntimeError::new_err(format!(
"Function '{}' has {} parameters, but only 0-5 parameters are supported",
name, param_count
)));
}
}
Ok(())
}
/// List registered Python functions
fn list_python_functions(&self) -> Vec<String> {
self.py_functions.keys().cloned().collect()
}
}
/// A Python module implemented in Rust using PyO3
#[pymodule]
fn rhai_python(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<RhaiEngine>()?;
m.add("__version__", "0.1.1")?;
Ok(())
}