fix rustpython

Author: marcomq

Cargo.lock

@@ -11,9 +11,8 @@ thiserror = "2.0"
 tokio = { version = "1.36.0", features = ["sync", "macros"] }
 
 [features]
-
-default = ["rustpython"]
-# default = ["pyo3"]
+default = ["pyo3"]
+# default = ["rustpython"]
 pyo3 = ["dep:pyo3"]
 rustpython = ["dep:rustpython-vm", "dep:rustpython-stdlib"]
 
@@ -27,9 +26,11 @@ tempfile = "3"
 [dependencies.rustpython-vm]
 version = "0.4.0"
 optional = true
+features = ["threading", "serde", "importlib"]
 [dependencies.rustpython-stdlib]
 version = "0.4.0"
 optional = true
+features = ["threading"]
 
 # The `full` feature for tokio is needed for the tests
 [dev-dependencies.tokio]

Cargo.toml

@@ -183,62 +183,34 @@ impl PyRunner {
     }
     /// Set python venv environment folder (does not change interpreter)
     pub async fn set_venv(&self, venv_path: &Path) -> Result<(), PyRunnerError> {
-        let set_venv_code = r##"
-import sys
-import os
-
-def add_venv_libs_to_syspath(venv_path):
-    """
-    Adds the site-packages folder (and .pth entries) from a virtual environment to sys.path.
-    
-    Args:
-        venv_path (str): Path to the root of the virtual environment.
-    """
-    if not os.path.isdir(venv_path):
-        raise ValueError(f"{venv_path} is not a directory")
-
-    if os.name == "nt":
-        # Windows: venv\Lib\site-packages
-        site_packages = os.path.join(venv_path, "Lib", "site-packages")
-    else:
-        # POSIX: venv/lib/pythonX.Y/site-packages
-        py_version = f"python{sys.version_info.major}.{sys.version_info.minor}"
-        site_packages = os.path.join(venv_path, "lib", py_version, "site-packages")
-
-    if not os.path.isdir(site_packages):
-        raise RuntimeError(f"Could not find site-packages in {venv_path}")
-
-    # Add site-packages itself
-    if site_packages not in sys.path:
-        sys.path.insert(0, site_packages)
-
-    # Process .pth files inside site-packages
-    for entry in os.listdir(site_packages):
-        if entry.endswith(".pth"):
-            pth_file = os.path.join(site_packages, entry)
-            try:
-                with open(pth_file, "r") as f:
-                    for line in f:
-                        line = line.strip()
-                        if not line or line.startswith("#"):
-                            continue
-                        if line.startswith("import "):
-                            # Execute import statements inside .pth files
-                            exec(line, globals(), locals())
-                        else:
-                            # Treat as a path
-                            extra_path = os.path.join(site_packages, line)
-                            if os.path.exists(extra_path) and extra_path not in sys.path:
-                                sys.path.insert(0, extra_path)
-            except Exception as e:
-                print(f"Warning: Could not process {pth_file}: {e}")
-
-    return site_packages
-"##;
+        if !venv_path.is_dir() {
+            return Err(PyRunnerError::PyError(format!(
+                "Could not find venv directory {}",
+                venv_path.display()
+            )));
+        }
+        let set_venv_code = include_str!("set_venv.py");
         self.run(&set_venv_code).await?;
+
+        let site_packages = if cfg!(target_os = "windows") {
+            venv_path.join("Lib").join("site-packages")
+        } else {
+            let version_code = "f\"python{sys.version_info.major}.{sys.version_info.minor}\"";
+            let py_version = self.eval(version_code).await?;
+            venv_path
+                .join("lib")
+                .join(py_version.as_str().unwrap())
+                .join("site-packages")
+        };
+        #[cfg(all(feature = "pyo3", not(feature = "rustpython")))]
+        let with_pth = "True";
+        #[cfg(feature = "rustpython")]
+        let with_pth = "False";
+
         self.run(&format!(
-            "add_venv_libs_to_syspath({})",
-            print_path_for_python(&venv_path.to_path_buf())
+            "add_venv_libs_to_syspath({}, {})",
+            print_path_for_python(&site_packages),
+            with_pth
         ))
         .await
     }

src/lib.rs

@@ -31,9 +31,7 @@ pub(crate) fn python_thread_main(mut receiver: mpsc::Receiver<PyCommand>) {
                     py.eval(&c_code, Some(&globals), None)
                         .and_then(|obj| py_any_to_json(py, &obj))
                 }
-                CmdType::RunFile(file) => {
-                    handle_run_file(py, &globals, file)
-                }
+                CmdType::RunFile(file) => handle_run_file(py, &globals, file),
                 CmdType::ReadVariable(var_name) => {
                     get_py_object(&globals, &var_name).and_then(|obj| py_any_to_json(py, &obj))
                 }
@@ -79,7 +77,7 @@ fn handle_run_file(
     globals: &pyo3::Bound<'_, PyDict>,
     file: std::path::PathBuf,
 ) -> PyResult<Value> {
-        let code = format!(
+    let code = format!(
         r#"
 import sys
 sys.path.insert(0, {})

src/pyo3_runner.rs

@@ -24,16 +24,22 @@ pub(crate) fn python_thread_main(mut receiver: mpsc::Receiver<PyCommand>) {
 
     interp.enter(|vm| {
         let scope = vm.new_scope_with_builtins();
+        vm.run_code_string(
+            scope.clone(),
+            "import sys; sys.path.append('./')",
+            "<init>".into(),
+        )
+        .unwrap();
         while let Some(cmd) = receiver.blocking_recv() {
             let result = match &cmd.cmd_type {
                 CmdType::RunCode(code) => vm
                     .run_code_string(scope.clone(), code, "<string, run>".to_owned())
                     .map(|_| Value::Null),
                 CmdType::EvalCode(code) => eval::eval(vm, code, scope.clone(), "<string, eval>")
                     .and_then(|obj| py_to_json(vm, &obj)),
-                CmdType::RunFile(file) => {
-                    handle_run_file(vm, scope.clone(), file).map(|_| Value::Null)
-                }
+                CmdType::RunFile(file) => vm
+                    .run_script(scope.clone(), file.to_str().unwrap())
+                    .map(|_| Value::Null),
                 CmdType::ReadVariable(var_name) => {
                     read_variable(vm, scope.clone(), var_name).and_then(|obj| py_to_json(vm, &obj))
                 }
@@ -55,16 +61,6 @@ pub(crate) fn python_thread_main(mut receiver: mpsc::Receiver<PyCommand>) {
     });
 }
 
-fn handle_run_file(
-    vm: &VirtualMachine,
-    scope: Scope,
-    file: &std::path::PathBuf,
-) -> PyResult<()> {
-    let dir = file.parent().unwrap().to_str().unwrap();
-    vm.insert_sys_path(vm.new_pyobj(dir))?;
-    vm.run_script(scope, file.to_str().unwrap())
-}
-
 fn read_variable(vm: &VirtualMachine, scope: Scope, var_name: &str) -> PyResult<PyObjectRef> {
     let parts: Vec<String> = var_name.split('.').map(|s| s.to_owned()).collect();
     let mut obj = scope.globals.get_item(&parts[0], vm)?;
@@ -161,6 +157,7 @@ fn py_to_json(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Value> {
         }
         return Ok(Value::Object(map));
     }
+    // fallback for types that can convert to primitives
     if let Some(b) = obj.clone().try_into_value::<bool>(vm).ok() {
         return Ok(Value::Bool(b));
     }
@@ -174,6 +171,7 @@ fn py_to_json(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Value> {
     }
 
     // Fallback for other types
+
     let fallback = obj.str(vm)?;
     Ok(Value::String(fallback.to_string()))
 }

src/rustpython_runner.rs

@@ -0,0 +1,38 @@
+import sys
+
+def add_venv_libs_to_syspath(site_packages, with_pth=False):
+    """
+    Adds the site-packages folder (and .pth entries) from a virtual environment to sys.path.
+    
+    Args:
+        venv_path (str): Path to the root of the virtual environment.
+    """
+
+    # Add site-packages itself
+    if site_packages not in sys.path:
+        sys.path.insert(0, site_packages)
+
+    # Process .pth files inside site-packages
+    if with_pth:
+        import os
+        for entry in os.listdir(site_packages):
+            if entry.endswith(".pth"):
+                pth_file = os.path.join(site_packages, entry)
+                try:
+                    with open(pth_file, "r") as f:
+                        for line in f:
+                            line = line.strip()
+                            if not line or line.startswith("#"):
+                                continue
+                            if line.startswith("import "):
+                                # Execute import statements inside .pth files
+                                exec(line, globals(), locals())
+                            else:
+                                # Treat as a path
+                                extra_path = os.path.join(site_packages, line)
+                                if extra_path not in sys.path:
+                                    sys.path.insert(0, extra_path)
+                except Exception as e:
+                    print(f"Warning: Could not process {pth_file}: {e}")
+
+    return site_packages