feat: remove dists site-packages's when python interpreter changes (p…

…refix-dev#896)

Co-authored-by: Bas Zalmstra <[email protected]>

src/environment.rs

@@ -206,7 +206,7 @@ pub async fn update_prefix_pypi(
     .await
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub enum PythonStatus {
     /// The python interpreter changed from `old` to `new`.
     Changed { old: PythonInfo, new: PythonInfo },

src/install_pypi.rs

@@ -23,6 +23,7 @@ use rattler_conda_types::{Platform, RepoDataRecord};
 use rattler_lock::{PypiPackageData, PypiPackageEnvironmentData};
 
 use std::collections::HashMap;
+use std::path::Path;
 use std::time::Duration;
 
 use uv_client::{FlatIndex, FlatIndexClient};
@@ -270,6 +271,37 @@ fn whats_the_plan<'a>(
     })
 }
 
+/// If the python interpreter is outdated, we need to uninstall all outdated site packages.
+/// from the old interpreter.
+async fn uninstall_outdated_site_packages(site_packages: &Path) -> miette::Result<()> {
+    // Check if the old interpreter is outdated
+    let mut installed = vec![];
+    for entry in std::fs::read_dir(site_packages).into_diagnostic()? {
+        let entry = entry.into_diagnostic()?;
+        if entry.file_type().into_diagnostic()?.is_dir() {
+            let path = entry.path();
+
+            let installed_dist = InstalledDist::try_from_path(&path);
+            let Ok(installed_dist) = installed_dist else {
+                continue;
+            };
+
+            if let Some(installed_dist) = installed_dist {
+                installed.push(installed_dist);
+            }
+        }
+    }
+
+    // Uninstall all packages in old site-packages
+    for dist_info in installed {
+        let _summary = uv_installer::uninstall(&dist_info)
+            .await
+            .expect("unistallation of old site-packages failed");
+    }
+
+    Ok(())
+}
+
 /// Installs and/or remove python distributions.
 // TODO: refactor arguments in struct
 #[allow(clippy::too_many_arguments)]
@@ -288,26 +320,21 @@ pub async fn update_python_distributions(
         return Ok(());
     };
 
-    let python_location = prefix.root().join(&python_info.path);
+    let platform = platform_host::Platform::current().expect("unsupported platform");
 
-    // Determine where packages would have been installed
-    let _python_version = (python_info.short_version.1 as u32, 0);
+    // If we have changed interpreter, we need to uninstall all site-packages from the old interpreter
+    if let PythonStatus::Changed { old, new: _ } = status {
+        let site_packages_path = prefix.root().join(&old.site_packages_path);
+        if site_packages_path.exists() {
+            uninstall_outdated_site_packages(&site_packages_path).await?;
+        }
+    }
 
+    // Determine the current environment markers.
     let python_record = conda_package
         .iter()
         .find(|r| is_python_record(r))
         .ok_or_else(|| miette::miette!("could not resolve pypi dependencies because no python interpreter is added to the dependencies of the project.\nMake sure to add a python interpreter to the [dependencies] section of the {PROJECT_MANIFEST}, or run:\n\n\tpixi add python"))?;
-
-    let platform = platform_host::Platform::current().expect("unsupported platform");
-    let interpreter =
-        Interpreter::query(&python_location, platform, &uv_context.cache).into_diagnostic()?;
-
-    tracing::debug!("[Install] Using Python Interpreter: {:?}", interpreter);
-
-    // Create a custom venv
-    let venv = PythonEnvironment::from_interpreter(interpreter, prefix.root());
-
-    // Determine the current environment markers.
     let tags = get_pypi_tags(
         Platform::current(),
         system_requirements,
@@ -331,6 +358,13 @@ pub async fn update_python_distributions(
     let in_memory_index = InMemoryIndex::default();
     let config_settings = ConfigSettings::default();
 
+    let python_location = prefix.root().join(&python_info.path);
+    let interpreter =
+        Interpreter::query(&python_location, platform, &uv_context.cache).into_diagnostic()?;
+
+    tracing::debug!("[Install] Using Python Interpreter: {:?}", interpreter);
+    // Create a custom venv
+    let venv = PythonEnvironment::from_interpreter(interpreter, prefix.root());
     // Prep the build context.
     let build_dispatch = BuildDispatch::new(
         &uv_context.registry_client,
@@ -438,7 +472,7 @@ pub async fn update_python_distributions(
         for dist_info in extraneous.iter().chain(reinstalls.iter()) {
             let summary = uv_installer::uninstall(dist_info)
                 .await
-                .expect("uinstall did not work");
+                .expect("uninstall did not work");
             tracing::debug!(
                 "Uninstalled {} ({} file{}, {} director{})",
                 dist_info.name(),

src/lock_file/update.rs

@@ -107,35 +107,35 @@ impl<'p> LockFileDerivedData<'p> {
             .unwrap_or_default();
         let pypi_records = self.pypi_records(environment, platform).unwrap_or_default();
 
-        let uv_context = match &self.uv_context {
-            None => {
-                let context = UvResolutionContext::from_project(self.project)?;
-                self.uv_context = Some(context.clone());
-                context
-            }
-            Some(context) => context.clone(),
-        };
-
-        // TODO(tim): get support for this somehow with uv
-        let env_variables = environment.project().get_env_variables(environment).await?;
+        if environment.has_pypi_dependencies() {
+            let uv_context = match &self.uv_context {
+                None => {
+                    let context = UvResolutionContext::from_project(self.project)?;
+                    self.uv_context = Some(context.clone());
+                    context
+                }
+                Some(context) => context.clone(),
+            };
 
-        // Update the prefix with Pypi records
-        environment::update_prefix_pypi(
-            environment.name(),
-            &prefix,
-            platform,
-            &repodata_records,
-            &pypi_records,
-            &python_status,
-            &environment.system_requirements(),
-            uv_context,
-            env_variables,
-        )
-        .await?;
+            let env_variables = environment.project().get_env_variables(environment).await?;
+            // Update the prefix with Pypi records
+            environment::update_prefix_pypi(
+                environment.name(),
+                &prefix,
+                platform,
+                &repodata_records,
+                &pypi_records,
+                &python_status,
+                &environment.system_requirements(),
+                uv_context,
+                env_variables,
+            )
+            .await?;
 
-        // Store that we updated the environment, so we won't have to do it again.
-        self.updated_pypi_prefixes
-            .insert(environment.clone(), prefix.clone());
+            // Store that we updated the environment, so we won't have to do it again.
+            self.updated_pypi_prefixes
+                .insert(environment.clone(), prefix.clone());
+        }
 
         Ok(prefix)
     }

tests/common/mod.rs

@@ -203,7 +203,7 @@ impl PixiControl {
         }
     }
 
-    /// Initialize pixi project inside a temporary directory. Returns a [`AddBuilder`]. To execute
+    /// Add dependencies to the project. Returns an [`AddBuilder`].
     /// the command and await the result call `.await` on the return value.
     pub fn add(&self, spec: &str) -> AddBuilder {
         AddBuilder {

tests/install_tests.rs

@@ -1,5 +1,7 @@
 mod common;
 
+use std::path::Path;
+
 use crate::common::builders::string_from_iter;
 use crate::common::package_database::{Package, PackageDatabase};
 use common::{LockFileExt, PixiControl};
@@ -8,6 +10,7 @@ use pixi::consts::DEFAULT_ENVIRONMENT_NAME;
 use rattler_conda_types::Platform;
 use serial_test::serial;
 use tempfile::TempDir;
+use uv_interpreter::PythonEnvironment;
 
 /// Should add a python version to the environment and lock file that matches the specified version
 /// and run it
@@ -194,3 +197,56 @@ async fn install_frozen() {
     assert_eq!(result.stdout.trim(), "Python 3.9.1");
     assert!(result.stderr.is_empty());
 }
+
+fn create_uv_environment(prefix: &Path, cache: &uv_cache::Cache) -> PythonEnvironment {
+    let python = if cfg!(target_os = "windows") {
+        prefix.join("python.exe")
+    } else {
+        prefix.join("bin/python")
+    };
+    let platform = platform_host::Platform::current().unwrap();
+    // Current interpreter and venv
+    let interpreter = uv_interpreter::Interpreter::query(&python, platform, &cache).unwrap();
+    uv_interpreter::PythonEnvironment::from_interpreter(interpreter, &prefix)
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+#[serial]
+#[cfg_attr(not(feature = "slow_integration_tests"), ignore)]
+async fn pypi_reinstall_python() {
+    let pixi = PixiControl::new().unwrap();
+    pixi.init().await.unwrap();
+    // Add and update lockfile with this version of python
+    pixi.add("python==3.11").with_install(true).await.unwrap();
+
+    // Add flask from pypi
+    pixi.add("flask")
+        .with_install(true)
+        .set_type(pixi::DependencyType::PypiDependency)
+        .await
+        .unwrap();
+
+    let prefix = pixi.project().unwrap().root().join(".pixi/envs/default");
+
+    let cache = uv_cache::Cache::temp().unwrap();
+
+    // Check if site-packages has entries
+    let env = create_uv_environment(&prefix, &cache);
+    let installed_311 = uv_installer::SitePackages::from_executable(&env).unwrap();
+    assert!(installed_311.iter().count() > 0);
+
+    // Reinstall python
+    pixi.add("python==3.12").with_install(true).await.unwrap();
+
+    // Check if site-packages has entries, should be empty now
+    let installed_311 = uv_installer::SitePackages::from_executable(&env).unwrap();
+
+    if cfg!(not(target_os = "windows")) {
+        // On non-windows the site-packages should be empty
+        assert!(installed_311.iter().count() == 0);
+    } else {
+        // Windows should still contain some packages
+        // This is because the site-packages is not prefixed with the python version
+        assert!(installed_311.iter().count() > 0);
+    }
+}