rustpy_ml/
runtime.rs

1use pyo3::prelude::*;
2use pyo3::types::PyModule;
3use pyo3::types::PyDict;
4use std::collections::HashMap;
5use once_cell::sync::OnceCell;
6use std::sync::Mutex;
7use crate::Error;
8
9static PY_INIT: OnceCell<()> = OnceCell::new();
10pub(crate) static GLOBALS: OnceCell<Mutex<HashMap<String, Py<PyAny>>>> = OnceCell::new();
11
12/// Initialize the Python runtime with optional ML library support
13/// This should be called before using any Python functionality
14pub fn init() -> crate::Result<()> {
15    PY_INIT.get_or_try_init(|| {
16        Python::with_gil(|py| {
17            // Set up Python path and environment
18            let _sys = py.import("sys")?;
19            
20            // Initialize global variables storage
21            GLOBALS.get_or_init(|| Mutex::new(HashMap::new()));
22            
23            Ok::<_, PyErr>(())
24        })?;
25
26        Ok::<_, Error>(())
27    })?;
28
29    Ok(())
30}
31
32/// Initialize with specific ML libraries
33pub fn init_with_ml(libraries: &[&str]) -> crate::Result<()> {
34    init()?;
35    
36    Python::with_gil(|py| {
37        for lib in libraries {
38            py.import(*lib).map_err(|e| {
39                Error::Interp(format!("Failed to import {}: {}", lib, e))
40            })?;
41        }
42        Ok(())
43    })
44}
45
46/// Execute Python code and return the result
47pub fn exec(code: &str) -> crate::Result<()> {
48    Python::with_gil(|py| {
49        PyModule::from_code_bound(py, code, "", "")?;
50        Ok(())
51    })
52}
53
54/// Evaluate Python expression and return the result
55pub fn eval<T>(code: &str) -> crate::Result<T>
56where
57    T: for<'py> FromPyObject<'py>,
58{
59    Python::with_gil(|py| {
60        // For evaluation, we create a module that computes the value
61        let wrapped_code = format!("__result__ = {}", code);
62        let module = PyModule::from_code_bound(py, &wrapped_code, "", "")?;
63        let result = module.getattr("__result__")?;
64        result.extract::<T>().map_err(Into::into)
65    })
66}
67
68/// Execute Python code with local variables
69pub fn exec_with_locals(code: &str, locals: &HashMap<String, Py<PyAny>>) -> crate::Result<()> {
70    Python::with_gil(|py| {
71        let globals = PyDict::new_bound(py);
72        for (key, value) in locals {
73            globals.set_item(key, value)?;
74        }
75        let module = PyModule::from_code_bound(py, code, "", "")?;
76        for (key, value) in locals {
77            module.add(key.as_str(), value)?;
78        }
79        Ok(())
80    })
81}
82
83/// Evaluate Python expression with local variables
84pub fn eval_with_locals<T>(code: &str, locals: &HashMap<String, Py<PyAny>>) -> crate::Result<T>
85where
86    T: for<'py> FromPyObject<'py>,
87{
88    Python::with_gil(|py| {
89        let wrapped_code = format!("__result__ = {}", code);
90        let module = PyModule::from_code_bound(py, &wrapped_code, "", "")?;
91        for (key, value) in locals {
92            module.add(key.as_str(), value)?;
93        }
94        let result = module.getattr("__result__")?;
95        result.extract::<T>().map_err(Into::into)
96    })
97}
98
99/// Import a Python module
100pub fn import(module_name: &str) -> crate::Result<Py<PyModule>> {
101    Python::with_gil(|py| {
102        let module = py.import(module_name)?;
103        Ok(module.unbind())
104    })
105}
106
107/// Call a Python function by module and function name
108pub fn call<T>(module_name: &str, func_name: &str, args: impl IntoPy<Py<PyAny>>) -> crate::Result<T>
109where
110    T: for<'py> FromPyObject<'py>,
111{
112    Python::with_gil(|py| {
113        let module = py.import(module_name)?;
114        let func = module.getattr(func_name)?;
115        let result = func.call1((args.into_py(py),))?;
116        result.extract::<T>().map_err(Into::into)
117    })
118}
119
120/// Store a global Python variable
121pub fn set_global(name: String, value: Py<PyAny>) -> crate::Result<()> {
122    let globals = GLOBALS.get().ok_or_else(|| {
123        Error::Interp("Runtime not initialized. Call init() first.".to_string())
124    })?;
125    
126    let mut globals = globals.lock().unwrap();
127    globals.insert(name, value);
128    Ok(())
129}
130
131/// Get a global Python variable
132pub fn get_global(name: &str) -> crate::Result<Option<Py<PyAny>>> {
133    let globals = GLOBALS.get().ok_or_else(|| {
134        Error::Interp("Runtime not initialized. Call init() first.".to_string())
135    })?;
136    
137    let globals = globals.lock().unwrap();
138    Ok(globals.get(name).map(|v| Python::with_gil(|py| v.clone_ref(py))))
139}
140
141/// Clear all global variables
142pub fn clear_globals() -> crate::Result<()> {
143    let globals = GLOBALS.get().ok_or_else(|| {
144        Error::Interp("Runtime not initialized. Call init() first.".to_string())
145    })?;
146    
147    let mut globals = globals.lock().unwrap();
148    globals.clear();
149    Ok(())
150}
rustpy_ml/runtime.rs

rustpy_ml/
runtime.rs
