From d95148f5151235ebada8bcd209af7059506a7bf1 Mon Sep 17 00:00:00 2001 From: Tim de Jager Date: Mon, 4 Mar 2024 11:56:01 +0100 Subject: [PATCH] feat: remove dists site-packages's when python interpreter changes (#896) Co-authored-by: Bas Zalmstra --- src/environment.rs | 2 +- src/install_pypi.rs | 64 +++++++++++++++++++++++++++++++---------- src/lock_file/update.rs | 54 +++++++++++++++++----------------- tests/common/mod.rs | 2 +- tests/install_tests.rs | 56 ++++++++++++++++++++++++++++++++++++ 5 files changed, 134 insertions(+), 44 deletions(-) diff --git a/src/environment.rs b/src/environment.rs index d648e3f0c..fd3e53167 100644 --- a/src/environment.rs +++ b/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 }, diff --git a/src/install_pypi.rs b/src/install_pypi.rs index 2bed5989b..e9363da1c 100644 --- a/src/install_pypi.rs +++ b/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(), diff --git a/src/lock_file/update.rs b/src/lock_file/update.rs index d07b33dee..c11a041f9 100644 --- a/src/lock_file/update.rs +++ b/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) } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index b6bea5f4d..3ed767226 100644 --- a/tests/common/mod.rs +++ b/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 { diff --git a/tests/install_tests.rs b/tests/install_tests.rs index 4fe73d72f..2ddd9bb1a 100644 --- a/tests/install_tests.rs +++ b/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); + } +}