| // Copyright 2022 The Pigweed Authors |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); you may not |
| // use this file except in compliance with the License. You may obtain a copy of |
| // the License at |
| // |
| // https://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT |
| // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| // License for the specific language governing permissions and limitations under |
| // the License. |
| |
| use std::sync::atomic::Ordering; |
| |
| use rustpython::{ |
| self, |
| vm::{ |
| self, identifier, |
| stdlib::{atexit, sys}, |
| PyObjectRef, |
| }, |
| InterpreterConfig, |
| }; |
| |
| pub mod ser; |
| |
| pub use ser::to_python_object; |
| |
| /// An ergonomic wrapper for a Python interpreter. |
| /// |
| /// This type is intentionally not exported outside of the crate for |
| /// now. Ideally all details of python code invocation will be private |
| /// to the crate. |
| pub(crate) struct Py { |
| interpreter: vm::Interpreter, |
| } |
| |
| impl Py { |
| /// Create a new Python interpreter. |
| pub fn new() -> Py { |
| let mut config = InterpreterConfig::new(); |
| config = config.init_stdlib(); |
| Py { |
| interpreter: config.interpreter(), |
| } |
| } |
| |
| /// Run Python code inside the interpreter. |
| /// |
| /// `init_scope` is run before the code is executed giving the caller the |
| /// ability to set up the scope before code execution. |
| pub fn run( |
| &mut self, |
| init_scope: impl FnOnce(&vm::VirtualMachine, &mut vm::scope::Scope) -> vm::PyResult<()>, |
| code: &str, |
| ) -> vm::PyResult<PyObjectRef> { |
| self.interpreter.enter(|vm| -> vm::PyResult<PyObjectRef> { |
| let mut scope = vm.new_scope_with_builtins(); |
| vm.import("site", None, 0)?; |
| |
| init_scope(vm, &mut scope)?; |
| |
| let result = vm |
| .compile(code, vm::compiler::Mode::Single, "<qg>".to_owned()) |
| .map_err(|err| vm.new_syntax_error(&err)) |
| .and_then(|code_obj| vm.run_code_obj(code_obj, scope.clone())); |
| Self::flush_std(vm); |
| result |
| }) |
| } |
| |
| fn flush_std(vm: &vm::VirtualMachine) { |
| if let Ok(stdout) = sys::get_stdout(vm) { |
| let _res = vm.call_method(&stdout, identifier!(vm, flush).as_str(), ()); |
| } |
| if let Ok(stderr) = sys::get_stderr(vm) { |
| let _res = vm.call_method(&stderr, identifier!(vm, flush).as_str(), ()); |
| } |
| } |
| } |
| |
| impl Drop for Py { |
| fn drop(&mut self) { |
| self.interpreter.enter(|vm| { |
| atexit::_run_exitfuncs(vm); |
| vm.state.finalizing.store(true, Ordering::Release); |
| |
| Self::flush_std(vm); |
| }); |
| } |
| } |
| |
| #[cfg(test)] |
| mod test { |
| use super::*; |
| use std::sync::{Arc, Mutex}; |
| |
| #[test] |
| fn function_callback() { |
| // The Python VM is treated as another thread context so types must |
| // implement Send and Sync. |
| let lines = Arc::new(Mutex::new(Vec::new())); |
| |
| // A clone for the python side |
| let py_lines = lines.clone(); |
| |
| let mut py = Py::new(); |
| let code = r#" |
| qg_test_func("line1") |
| qg_test_func("line2") |
| "#; |
| py.run( |
| // Scope initialization function. |
| |vm, scope| { |
| // Add a `rust_print` function to the scope before running the code. |
| scope.globals.set_item( |
| "qg_test_func", |
| vm.ctx |
| .new_function("qg_test_func", move |s: String| { |
| println!("qg_test_func: {}", s); |
| py_lines.lock().unwrap().push(s); |
| }) |
| .into(), |
| vm, |
| )?; |
| |
| Ok(()) |
| }, |
| code, |
| ) |
| .unwrap(); |
| |
| assert_eq!( |
| *lines.lock().unwrap(), |
| vec!["line1".to_string(), "line2".to_string()] |
| ); |
| } |
| } |