qg/src/py/mod.rs

// 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,
};

/// 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()]
        );
    }
}