wild_doc_script_python/
lib.rs

1use std::{ffi::CString, path::PathBuf, sync::Arc};
2
3use parking_lot::Mutex;
4use pyo3::{
5    pyfunction,
6    types::{PyCapsule, PyDict, PyModule},
7    wrap_pyfunction, PyObject, PyResult, Python,
8};
9use wild_doc_script::{
10    anyhow::Result, async_trait, IncludeAdaptor, Stack, WildDocScript, WildDocValue,
11};
12
13pub struct WdPy {}
14
15#[async_trait(?Send)]
16impl<I: IncludeAdaptor + Send> WildDocScript<I> for WdPy {
17    fn new(_: Arc<Mutex<I>>, _: PathBuf, _: &Stack) -> Result<Self> {
18        let _ = Python::with_gil(|py| -> PyResult<()> {
19            let builtins = PyModule::import(py, "builtins")?;
20
21            let wd = PyModule::new(py, "wd")?;
22            wd.add_function(wrap_pyfunction!(wdv, wd)?)?;
23
24            builtins.add_function(wrap_pyfunction!(wdv, builtins)?)?;
25
26            builtins.add_submodule(wd)?;
27
28            Ok(())
29        });
30        Ok(WdPy {})
31    }
32
33    async fn evaluate_module(&mut self, _: &str, code: &str, stack: &Stack) -> Result<()> {
34        Python::with_gil(|py| -> PyResult<()> {
35            let builtins = PyModule::import(py, "builtins")?;
36            builtins.set_item(
37                "wdvars",
38                PyCapsule::new(py, stack.clone(), Some(CString::new("builtins.wdstack")?))?,
39            )?;
40
41            py.run(code, None, None)
42        })?;
43        Ok(())
44    }
45
46    async fn eval(&mut self, code: &str, stack: &Stack) -> Result<WildDocValue> {
47        Ok(WildDocValue::Binary(
48            Python::with_gil(|py| -> PyResult<PyObject> {
49                let builtins = PyModule::import(py, "builtins")?;
50                builtins.set_item(
51                    "wdvars",
52                    PyCapsule::new(py, stack.clone(), Some(CString::new("builtins.wdstack")?))?,
53                )?;
54                py.eval(code, None, None)?.extract()
55            })?
56            .to_string()
57            .into_bytes(),
58        ))
59    }
60}
61
62#[pyfunction]
63#[pyo3(name = "v")]
64fn wdv(_py: Python, key: String) -> PyResult<PyObject> {
65    Python::with_gil(|py| -> PyResult<PyObject> {
66        let stack: &Stack =
67            unsafe { PyCapsule::import(py, CString::new("builtins.wdstack")?.as_ref())? };
68
69        if let Some(v) = stack.get(&Arc::new(key.into())) {
70            return PyModule::from_code(
71                py,
72                r#"
73import json
74
75def v(data):
76    return json.loads(data)
77"#,
78                "",
79                "",
80            )?
81            .getattr("v")?
82            .call1((v.to_string(),))?
83            .extract();
84        }
85
86        Ok(PyDict::new(py).into())
87    })
88}