neutralts/bif/
exec_python.rs

1// pyo3 = { version = "0.26.2", features = [] }
2
3use crate::{bif::BifError, Value};
4use pyo3::prelude::*;
5use pyo3::types::{PyList, PyModule};
6use std::env;
7use std::path::Path;
8use std::process::Command;
9
10pub struct PythonExecutor;
11
12impl PythonExecutor {
13    pub(crate) fn exec_py(
14        file: &str,
15        params_value: &Value,
16        callback_name: &str,
17        schema: Option<&Value>,
18        schema_data: Option<&Value>,
19        venv_path: Option<&str>,
20    ) -> Result<Value, BifError> {
21        if let Some(venv) = venv_path {
22            Self::setup_venv(venv)?;
23        }
24
25        Python::initialize();
26
27        Python::attach(|py| -> PyResult<Value> {
28            let params = Self::prepare_python_params(py, params_value)?;
29            Self::setup_python_path(py, file)?;
30            Self::execute_python_callback(py, file, callback_name, params, schema, schema_data)
31        })
32        .map_err(|e| BifError {
33            msg: format!(
34                "Error executing callback function '{}': {}",
35                callback_name, e
36            ),
37            name: "python_callback".to_string(),
38            file: file.to_string(),
39            src: file.to_string(),
40        })
41    }
42
43    fn setup_venv(venv_path: &str) -> Result<(), BifError> {
44        let path = Path::new(venv_path);
45        if !path.exists() {
46            return Err(BifError {
47                msg: format!("Venv path '{}' does not exist", venv_path),
48                name: "venv_error".to_string(),
49                file: "".to_string(),
50                src: "".to_string(),
51            });
52        }
53
54        let python_executable = if cfg!(unix) {
55            format!("{}/bin/python", venv_path)
56        } else {
57            format!("{}\\Scripts\\python.exe", venv_path)
58        };
59
60        if !Path::new(&python_executable).exists() {
61            return Err(BifError {
62                msg: format!("Python executable not found: {}", python_executable),
63                name: "venv_error".to_string(),
64                file: "".to_string(),
65                src: "".to_string(),
66            });
67        }
68
69        env::set_var("PYTHON_EXECUTABLE", &python_executable);
70        env::set_var("VIRTUAL_ENV", venv_path);
71
72        let output = Command::new(&python_executable)
73            .arg("-c")
74            .arg("import sys; print(sys.prefix); print(':'.join(sys.path))")
75            .output()
76            .map_err(|e| BifError {
77                msg: format!("Failed to get Python path info: {}", e),
78                name: "venv_error".to_string(),
79                file: "".to_string(),
80                src: "".to_string(),
81            })?;
82
83        if output.status.success() {
84            let output_str = String::from_utf8_lossy(&output.stdout);
85            let lines: Vec<&str> = output_str.trim().split('\n').collect();
86            if lines.len() >= 2 {
87                env::set_var("PYTHONHOME", lines[0]);
88                env::set_var("PYTHONPATH", lines[1]);
89            }
90        }
91
92        Ok(())
93    }
94
95    fn prepare_python_params<'py>(py: Python<'py>, params_value: &Value) -> PyResult<Py<PyAny>> {
96        let params_json = serde_json::to_string(params_value).map_err(|e| {
97            pyo3::exceptions::PyValueError::new_err(format!("Failed to serialize params: {}", e))
98        })?;
99        let json_mod = PyModule::import(py, "json")?;
100        let loads = json_mod.getattr("loads")?;
101        let py_obj = loads.call1((params_json,))?;
102        let py_object: Py<PyAny> = py_obj.extract()?;
103        Ok(py_object)
104    }
105
106    fn setup_python_path(py: Python, file: &str) -> PyResult<()> {
107        let dir_path = Path::new(file).parent().unwrap_or_else(|| Path::new("."));
108        let sys = PyModule::import(py, "sys")?;
109        let path_attr = sys.getattr("path")?;
110        let path = path_attr.cast::<PyList>()?;
111        if let Some(dir_str) = dir_path.to_str() {
112            path.append(dir_str)?;
113        } else {
114            return Err(pyo3::exceptions::PyValueError::new_err(
115                "Invalid directory path encoding",
116            ));
117        }
118        Ok(())
119    }
120
121    fn execute_python_callback<'py>(
122        py: Python<'py>,
123        file: &str,
124        callback_name: &str,
125        params: Py<PyAny>,
126        schema: Option<&Value>,
127        schema_data: Option<&Value>,
128    ) -> PyResult<Value> {
129        let module_name = Self::extract_module_name(file)?;
130        let module = PyModule::import(py, &module_name)?;
131
132        // https://github.com/FranBarInstance/neutralts/issues/2
133        if module.hasattr("__NEUTRAL_SCHEMA__")? {
134            module.delattr("__NEUTRAL_SCHEMA__")?;
135        }
136
137        if let Some(schema_value) = schema {
138            let schema_py = Self::prepare_python_params(py, schema_value)?;
139            module.setattr("__NEUTRAL_SCHEMA__", schema_py)?;
140        }
141
142        if module.hasattr("__NEUTRAL_SCHEMA_DATA__")? {
143            module.delattr("__NEUTRAL_SCHEMA_DATA__")?;
144        }
145
146        if let Some(schema_data_value) = schema_data {
147            let schema_data_py = Self::prepare_python_params(py, schema_data_value)?;
148            module.setattr("__NEUTRAL_SCHEMA_DATA__", schema_data_py)?;
149        }
150
151        let callback_func = module.getattr(callback_name).map_err(|_| {
152            pyo3::exceptions::PyAttributeError::new_err(format!(
153                "Module does not have function '{}'",
154                callback_name
155            ))
156        })?;
157        let result_any = callback_func.call1((params,))?;
158        let result_obj: Py<PyAny> = result_any.extract()?;
159        Self::convert_python_result_to_json(py, result_obj)
160    }
161
162    fn extract_module_name(file: &str) -> PyResult<String> {
163        Path::new(file)
164            .file_stem()
165            .and_then(|s| s.to_str())
166            .map(|s| s.to_string())
167            .ok_or_else(|| {
168                pyo3::exceptions::PyValueError::new_err(
169                    "Could not extract module name from file path",
170                )
171            })
172    }
173
174    fn convert_python_result_to_json<'py>(py: Python<'py>, result: Py<PyAny>) -> PyResult<Value> {
175        let json_module = PyModule::import(py, "json")?;
176        let json_dumps = json_module.getattr("dumps")?;
177        let json_string: String = json_dumps.call1((result,))?.extract()?;
178        serde_json::from_str(&json_string).map_err(|e| {
179            pyo3::exceptions::PyValueError::new_err(format!("Error parsing JSON: {}", e))
180        })
181    }
182}