fix: improve support for rust interpreter python imports from venv

Signed-off-by: Nick Mitchell <[email protected]>

Author: starpit

.github/workflows/rust-interpreter.yml

@@ -31,4 +31,8 @@ jobs:
             (curl -fsSL https://ollama.com/install.sh | sudo -E sh && sleep 2)
           wait
       - name: Run interpreter tests
-        run: npm run test:interpreter
+        run: |
+          python3.12 -mvenv venv
+          source venv/bin/activate
+          pip install nested-diff
+          npm run test:interpreter

pdl-live-react/src-tauri/Cargo.lock

@@ -38,7 +38,7 @@ minijinja = { version = "2.9.0", features = ["custom_syntax"] }
 #ollama-rs = { version = "0.3.0", features = ["stream"] }
 ollama-rs = { git = "https://github.com/starpit/ollama-rs.git", branch = "tools-pub-7", features = ["stream"] }
 owo-colors = "4.2.0"
-rustpython-vm = { git="https://github.com/RustPython/RustPython.git" } # "0.4.0"
+rustpython-vm = { git="https://github.com/RustPython/RustPython.git", features= ["importlib", "threading", "encodings"] } # "0.4.0"
 async-recursion = "1.1.1"
 tokio-stream = "0.1.17"
 tokio = { version = "1.44.1", features = ["io-std"] }

pdl-live-react/src-tauri/Cargo.toml

@@ -759,28 +759,55 @@ impl<'a> Interpreter<'a> {
         _state: &mut State,
     ) -> BodyInterpretation {
         use rustpython_vm as vm;
-        let interp = vm::Interpreter::with_init(vm::Settings::default(), |vm| {
-            //vm.add_native_modules(rustpython_stdlib::get_module_inits());
+
+        let mut settings = rustpython_vm::Settings::default();
+
+        // add PYTHONPATH to sys.path
+        settings.path_list.extend(get_paths("PDLPYTHONPATH"));
+        settings.path_list.extend(get_paths("PYTHONPATH"));
+
+        if let Ok(venv) = ::std::env::var("VIRTUAL_ENV") {
+            let path = ::std::path::PathBuf::from(venv).join(if cfg!(windows) {
+                "lib/site-packages"
+            } else {
+                // TODO generalize this!
+                "lib/python3.12/site-packages"
+            });
+            settings = settings.with_path(path.display().to_string());
+        }
+
+        let interp = vm::Interpreter::with_init(settings, |vm| {
+            vm.add_native_modules(rustpython_stdlib::get_module_inits());
             vm.add_frozen(rustpython_pylib::FROZEN_STDLIB);
         });
         interp.enter(|vm| -> BodyInterpretation {
             let scope = vm.new_scope_with_builtins();
 
-            // TODO vm.new_syntax_error(&err, Some(block.code.as_str()))
-            let code_obj = match vm.compile(
-                block.code.as_str(),
-                vm::compiler::Mode::Exec,
+            // Sigh, this is copy-pasted from RustPython/src/lib.rs
+            // `run_rustpython` as of 20250416 commit hash
+            // a917da3b1. Without this (and also: importlib and
+            // encodings features on rustpython-vm crate), then
+            // pulling in venvs does not work.
+            match vm.run_code_string(
+                vm.new_scope_with_builtins(),
+                "import sys; sys.path.insert(0, '')",
                 "<embedded>".to_owned(),
             ) {
-                Ok(x) => Ok(x),
-                Err(exc) => Err(PdlError::from(format!(
-                    "Syntax error in Python code {:?}",
-                    exc
-                ))),
+                Ok(_) => Ok(()),
+                Err(exc) => {
+                    vm.print_exception(exc);
+                    Err(PdlError::from("Error setting up Python site path"))
+                }
             }?;
+            let site_result = vm.import("site", 0);
+            if site_result.is_err() {
+                println!(
+                    "Failed to import site, consider adding the Lib directory to your RUSTPYTHONPATH \
+                     environment variable",
+                );
+            }
 
-            // TODO vm.print_exception(exc);
-            match vm.run_code_obj(code_obj, scope.clone()) {
+            match vm.run_code_string(scope.clone(), block.code.as_str(), "<embedded>".to_owned()) {
                 Ok(_) => Ok(()),
                 Err(exc) => {
                     vm.print_exception(exc);
@@ -1500,3 +1527,30 @@ pub fn load_scope(
 
     Ok(scope)
 }
+
+/// Helper function to retrieve a sequence of paths from an environment variable.
+fn get_paths(env_variable_name: &str) -> impl Iterator<Item = String> + '_ {
+    ::std::env::var_os(env_variable_name)
+        .into_iter()
+        .flat_map(move |paths| {
+            split_paths(&paths)
+                .map(|path| {
+                    path.into_os_string()
+                        .into_string()
+                        .unwrap_or_else(|_| panic!("{env_variable_name} isn't valid unicode"))
+                })
+                .collect::<Vec<_>>()
+        })
+}
+
+#[cfg(not(target_os = "wasi"))]
+pub(crate) use ::std::env::split_paths;
+#[cfg(target_os = "wasi")]
+pub(crate) fn split_paths<T: AsRef<std::ffi::OsStr> + ?Sized>(
+    s: &T,
+) -> impl Iterator<Item = std::path::PathBuf> + '_ {
+    use std::os::wasi::ffi::OsStrExt;
+    let s = s.as_ref().as_bytes();
+    s.split(|b| *b == b':')
+        .map(|x| std::ffi::OsStr::from_bytes(x).to_owned().into())
+}

pdl-live-react/src-tauri/src/pdl/interpreter.rs

@@ -318,6 +318,45 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn text_python_two_code_result_dict() -> Result<(), Box<dyn Error>> {
+        let program = json!({
+            "text": [
+                { "lang": "python",
+                   "code":"print('hi ho'); result = {\"foo\": 3}"
+                },
+                { "lang": "python",
+                   "code":"import os; print('hi ho'); result = {\"foo\": 4}"
+                }
+            ]
+        });
+
+        let (_, messages, _) = run_json(program, streaming(), initial_scope())?;
+        assert_eq!(messages.len(), 2);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[0].content, "{'foo': 3}");
+        assert_eq!(messages[1].role, MessageRole::User);
+        assert_eq!(messages[1].content, "{'foo': 4}");
+        Ok(())
+    }
+
+    // TODO: illegal instruction, but only during tests
+    #[test]
+    fn text_python_code_import_venv() -> Result<(), Box<dyn Error>> {
+        let program = json!({
+            "include": "./tests/cli/code-python.pdl"
+        });
+
+        let (_, messages, _) = run_json(program, streaming(), initial_scope())?;
+        assert_eq!(messages.len(), 1);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(
+            messages[0].content,
+            "{'foo': None, 'diff': {'D': {'two': {'N': 3}}}}"
+        );
+        Ok(())
+    }
+
     #[test]
     fn text_read_file_text() -> Result<(), Box<dyn Error>> {
         let program = json!({

pdl-live-react/src-tauri/src/pdl/interpreter_tests.rs

@@ -1,5 +1,10 @@
 lang: python
 code: |
-  import os
-  print(f"hi ho {os.getcwd()}")
-  result = {"foo": 3}
+  import sys # test import stdlib
+  import os # test import stdlib
+  print(f"!!! {sys.path}")
+  #import textdistance # test import from venv
+  from nested_diff import diff # test import from venv
+  a = {'one': 1, 'two': 2, 'three': 3}
+  b = {'one': 1, 'two': 3, 'three': 3}
+  result = {"foo": os.getenv("FOO999999999999999999999999"), "diff": diff(a, b, O=False, U=False)}