rustpy_ml/

macro_embed.rs

use pyo3::prelude::*;
use pyo3::types::PyModule;
use std::collections::HashMap;

/// Execute Python code as a string and return the result
/// This is the runtime function called by the python! macro
pub fn run_python_code<T>(code: &str, locals: Option<HashMap<String, Py<PyAny>>>) -> crate::Result<T>
where
    T: for<'py> FromPyObject<'py>,
{
    crate::runtime::init()?;
    
    Python::with_gil(|py| {
        // Always create a locals dict for consistent behavior
        let locals_dict = pyo3::types::PyDict::new_bound(py);
        
        // Add globals first
        let globals = crate::runtime::GLOBALS.get()
            .ok_or_else(|| crate::Error::Interp("Runtime not initialized".to_string()))?;
        let globals_map = globals.lock().unwrap();
        for (key, value) in globals_map.iter() {
            locals_dict.set_item(key, value)?;
        }
        drop(globals_map); // Release the lock
        
        // Track initial keys
        let mut initial_keys = std::collections::HashSet::new();
        for key in locals_dict.keys() {
            if let Ok(k) = key.extract::<String>() {
                initial_keys.insert(k);
            }
        }
        
        if let Some(locals_map) = &locals {
            for (key, value) in locals_map {
                initial_keys.insert(key.clone());
                locals_dict.set_item(key, value)?;
            }
        }
        
        // Try to evaluate as expression first (wrapping in __result__)
        let wrapped_code = format!("__result__ = {}", code.trim());
        
        let result = py.run_bound(&wrapped_code, None, Some(&locals_dict));
            
        match result {
            Ok(_) => {
                // Successfully evaluated as expression, get the result
                let result_obj = locals_dict.get_item("__result__")?.ok_or_else(|| {
                    pyo3::exceptions::PyKeyError::new_err("__result__ not found")
                })?;
                result_obj.extract::<T>().map_err(Into::into)
            },
            Err(_) => {
                // Not an expression - try to execute as statements
                py.run_bound(code, None, Some(&locals_dict))?;
                
                // After execution, look for 'result' variable first (convention)
                if let Some(result) = locals_dict.get_item("result")? {
                    // Update globals before returning
                    let mut globals_map = globals.lock().unwrap();
                    for (key, value) in locals_dict.iter() {
                        if let Ok(k) = key.extract::<String>() {
                            if !k.starts_with("__") && !k.ends_with("__") && !initial_keys.contains(&k) {
                                globals_map.insert(k, value.into());
                            }
                        }
                    }
                    return result.extract::<T>().map_err(Into::into);
                }
                
                // Look for the last non-dunder, non-module variable added during execution
                let items: Vec<(String, pyo3::Bound<pyo3::PyAny>)> = locals_dict.items().extract()?;
                for (key, value) in items.iter().rev() {
                    // Skip builtins, private, and initial keys
                    if key.starts_with("__") || key.starts_with("_") || initial_keys.contains(key) {
                        continue;
                    }
                    
                    // Skip modules (they have __name__ attribute)
                    if value.hasattr("__name__").unwrap_or(false) && 
                       value.hasattr("__package__").unwrap_or(false) {
                        continue;
                    }
                    
                    // Update globals before returning
                    let mut globals_map = globals.lock().unwrap();
                    for (k, v) in locals_dict.iter() {
                        if let Ok(k_str) = k.extract::<String>() {
                            if !k_str.starts_with("__") && !k_str.ends_with("__") && !initial_keys.contains(&k_str) {
                                globals_map.insert(k_str, v.into());
                            }
                        }
                    }
                    return value.extract::<T>().map_err(Into::into);
                }
                
                // Return None if no result found
                py.None().extract::<T>(py).map_err(Into::into)
            }
        }
    })
}

/// Execute Python code without returning a value
pub fn exec_python_code(code: &str, locals: Option<HashMap<String, Py<PyAny>>>) -> crate::Result<()> {
    crate::runtime::init()?;
    
    Python::with_gil(|py| {
        // Create locals dictionary
        let locals_dict = pyo3::types::PyDict::new_bound(py);
        
        // Add globals first
        let globals = crate::runtime::GLOBALS.get()
            .ok_or_else(|| crate::Error::Interp("Runtime not initialized".to_string()))?;
        let globals_map = globals.lock().unwrap();
        for (key, value) in globals_map.iter() {
            locals_dict.set_item(key, value)?;
        }
        drop(globals_map); // Release the lock
        
        // Add provided locals to the dictionary (can override globals)
        if let Some(locals_map) = locals {
            for (key, value) in locals_map {
                locals_dict.set_item(key, value)?;
            }
        }
        
        // Execute code with locals
        py.run_bound(code, None, Some(&locals_dict))?;
        
        // Update globals with any new variables defined in locals
        let mut globals_map = globals.lock().unwrap();
        for (key, value) in locals_dict.iter() {
            // Skip dunders and built-ins
            if let Ok(key_str) = key.extract::<String>() {
                if !key_str.starts_with("__") && !key_str.ends_with("__") {
                    globals_map.insert(key_str, value.into());
                }
            }
        }
        
        Ok(())
    })
}

/// Simple macro to run Python code
/// Usage: python!("print('Hello from Python')")
#[macro_export]
macro_rules! python {
    // Simple case: just code, return PyObject
    ($code:expr) => {{
        $crate::macro_embed::run_python_code::<pyo3::PyObject>($code, None)
    }};
    
    // With type annotation: python!(-> i32, "1 + 1")
    (-> $ty:ty, $code:expr) => {{
        $crate::macro_embed::run_python_code::<$ty>($code, None)
    }};
    
    // With locals: python!({"x" => 5}, "x * 2")
    ($locals:expr, $code:expr) => {{
        $crate::macro_embed::run_python_code::<pyo3::PyObject>($code, Some($locals))
    }};
    
    // With locals and type: python!({"x" => 5}, -> i32, "x * 2")
    ($locals:expr, -> $ty:ty, $code:expr) => {{
        $crate::macro_embed::run_python_code::<$ty>($code, Some($locals))
    }};
}

/// Macro to execute Python code without returning a value
/// Usage: py_exec!("import numpy as np\nprint(np.__version__)")
#[macro_export]
macro_rules! py_exec {
    ($code:expr) => {{
        $crate::macro_embed::exec_python_code($code, None)
    }};
    
    ($locals:expr, $code:expr) => {{
        $crate::macro_embed::exec_python_code($code, Some($locals))
    }};
}

/// Macro to evaluate Python expression with type inference
/// Usage: let result: i32 = py_eval!("1 + 1");
#[macro_export]
macro_rules! py_eval {
    ($code:expr) => {{
        $crate::runtime::eval($code)
    }};
}

/// Helper macro to create a locals HashMap
/// Usage: py_locals!{"x" => 5, "y" => 10}
#[macro_export]
macro_rules! py_locals {
    ($($key:expr => $value:expr),* $(,)?) => {{
        let mut map = std::collections::HashMap::new();
        pyo3::Python::with_gil(|py| {
            $(
                map.insert($key.to_string(), pyo3::IntoPy::into_py($value, py));
            )*
        });
        map
    }};
}