feat: add reinstallation method to installer and transaction (#1128)

tdejager · web-flow · commit c7e2197d0bf8 · 2025-03-07T15:03:48.000+01:00
diff --git a/crates/rattler-bin/src/commands/create.rs b/crates/rattler-bin/src/commands/create.rs
@@ -288,6 +288,7 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
         let transaction = Transaction::from_current_and_desired(
             installed_packages,
             required_packages,
+            None,
             install_platform,
         )?;
 
diff --git a/crates/rattler/src/install/installer/mod.rs b/crates/rattler/src/install/installer/mod.rs
@@ -3,7 +3,7 @@ mod error;
 mod indicatif;
 mod reporter;
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     future::ready,
     path::{Path, PathBuf},
     sync::Arc,
@@ -20,7 +20,7 @@ use itertools::Itertools;
 use rattler_cache::package_cache::{CacheLock, CacheReporter};
 use rattler_conda_types::{
     prefix_record::{Link, LinkType},
-    Platform, PrefixRecord, RepoDataRecord,
+    PackageName, Platform, PrefixRecord, RepoDataRecord,
 };
 use rattler_networking::retry_policies::default_retry_policy;
 pub use reporter::Reporter;
@@ -50,6 +50,7 @@ pub struct Installer {
     target_platform: Option<Platform>,
     apple_code_sign_behavior: AppleCodeSignBehavior,
     alternative_target_prefix: Option<PathBuf>,
+    reinstall_packages: Option<HashSet<PackageName>>,
     // TODO: Determine upfront if these are possible.
     // allow_symbolic_links: Option<bool>,
     // allow_hard_links: Option<bool>,
@@ -216,10 +217,27 @@ impl Installer {
         }
     }
 
+    /// Set the packages that we want explicitly to be reinstalled.
+    #[must_use]
+    pub fn with_reinstall_packages(self, reinstall: HashSet<PackageName>) -> Self {
+        Self {
+            reinstall_packages: Some(reinstall),
+            ..self
+        }
+    }
+
+    /// Set the packages that we want explicitly to be reinstalled.
+    /// This function is similar to [`Self::with_reinstall_packages`],but
+    /// modifies an existing instance.
+    pub fn set_reinstall_packages(&mut self, reinstall: HashSet<PackageName>) -> &mut Self {
+        self.reinstall_packages = Some(reinstall);
+        self
+    }
+
     /// Sets the packages that are currently installed in the prefix. If this
     /// is not set, the installation process will first figure this out.
     ///
-    /// This function is similar to [`Self::set_installed_packages`],but
+    /// This function is similar to [`Self::with_installed_packages`],but
     /// modifies an existing instance.
     pub fn set_installed_packages(&mut self, installed: Vec<PrefixRecord>) -> &mut Self {
         self.installed = Some(installed);
@@ -314,6 +332,7 @@ impl Installer {
         let transaction = Transaction::from_current_and_desired(
             installed,
             records.into_iter().collect::<Vec<_>>(),
+            self.reinstall_packages,
             target_platform,
         )?;
 
diff --git a/crates/rattler/src/install/test_utils.rs b/crates/rattler/src/install/test_utils.rs
@@ -1,16 +1,21 @@
-use std::path::{Path, PathBuf};
+use std::{
+    path::{Path, PathBuf},
+    str::FromStr,
+};
 
 use futures::TryFutureExt;
-use rattler_conda_types::{PrefixRecord, RepoDataRecord};
+use rattler_conda_types::{Platform, PrefixRecord, RepoDataRecord, Version};
 use rattler_networking::retry_policies::default_retry_policy;
 use transaction::{Transaction, TransactionOperation};
+use url::Url;
 
 use crate::{
+    get_repodata_record,
     install::{transaction, unlink_package, InstallDriver, InstallOptions},
     package_cache::PackageCache,
 };
 
-use super::driver::PostProcessResult;
+use super::{driver::PostProcessResult, link_package, PythonInfo};
 
 /// Install a package into the environment and write a `conda-meta` file that
 /// contains information about how the file was linked.
@@ -157,3 +162,42 @@ pub fn find_prefix_record<'a>(
         .iter()
         .find(|r| r.repodata_record.package_record.name.as_normalized() == name)
 }
+
+pub async fn download_and_get_prefix_record(
+    target_prefix: &Path,
+    package_url: Url,
+    sha256_hash: &str,
+) -> PrefixRecord {
+    let package_path = tools::download_and_cache_file_async(package_url, sha256_hash)
+        .await
+        .unwrap();
+
+    let package_dir = tempfile::TempDir::new().unwrap();
+
+    // Create package cache
+    rattler_package_streaming::fs::extract(&package_path, package_dir.path()).unwrap();
+
+    let py_info =
+        PythonInfo::from_version(&Version::from_str("3.10").unwrap(), None, Platform::Linux64)
+            .unwrap();
+    let install_options = InstallOptions {
+        python_info: Some(py_info),
+        ..InstallOptions::default()
+    };
+
+    let install_driver = InstallDriver::default();
+    // Link the package
+    let paths = link_package(
+        package_dir.path(),
+        target_prefix,
+        &install_driver,
+        install_options,
+    )
+    .await
+    .unwrap();
+
+    let repodata_record = get_repodata_record(&package_path);
+    // Construct a PrefixRecord for the package
+
+    PrefixRecord::from_repodata_record(repodata_record, None, None, paths, None, None)
+}
diff --git a/crates/rattler/src/install/transaction.rs b/crates/rattler/src/install/transaction.rs
@@ -1,6 +1,6 @@
 use std::collections::HashSet;
 
-use rattler_conda_types::{PackageRecord, Platform};
+use rattler_conda_types::{PackageName, PackageRecord, Platform};
 
 use crate::install::{python::PythonInfoError, PythonInfo};
 
@@ -125,13 +125,15 @@ impl<Old: AsRef<New>, New> Transaction<Old, New> {
 
 impl<Old: AsRef<PackageRecord>, New: AsRef<PackageRecord>> Transaction<Old, New> {
     /// Constructs a [`Transaction`] by taking the current situation and diffing
-    /// that against the desired situation.
+    /// that against the desired situation. You can specify a set of package names
+    /// that should be reinstalled even if their content has not changed.
     pub fn from_current_and_desired<
         CurIter: IntoIterator<Item = Old>,
         NewIter: IntoIterator<Item = New>,
     >(
         current: CurIter,
         desired: NewIter,
+        reinstall: Option<HashSet<PackageName>>,
         platform: Platform,
     ) -> Result<Self, TransactionError>
     where
@@ -150,6 +152,7 @@ impl<Old: AsRef<PackageRecord>, New: AsRef<PackageRecord>> Transaction<Old, New>
         };
 
         let mut operations = Vec::new();
+        let reinstall = reinstall.unwrap_or_default();
 
         let mut current_map = current_iter
             .clone()
@@ -179,7 +182,9 @@ impl<Old: AsRef<PackageRecord>, New: AsRef<PackageRecord>> Transaction<Old, New>
             let old_record = current_map.remove(name);
 
             if let Some(old_record) = old_record {
-                if !describe_same_content(record.as_ref(), old_record.as_ref()) {
+                if !describe_same_content(record.as_ref(), old_record.as_ref())
+                    || reinstall.contains(&record.as_ref().name)
+                {
                     // if the content changed, we need to reinstall (remove and install)
                     operations.push(TransactionOperation::Change {
                         old: old_record,
@@ -249,3 +254,42 @@ fn describe_same_content(from: &PackageRecord, to: &PackageRecord) -> bool {
     // Otherwise, just check that the name, version and build string match
     from.name == to.name && from.version == to.version && from.build == to.build
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+
+    use rattler_conda_types::Platform;
+
+    use crate::install::{
+        test_utils::download_and_get_prefix_record, Transaction, TransactionOperation,
+    };
+    use assert_matches::assert_matches;
+
+    #[tokio::test]
+    async fn test_reinstall_package() {
+        let environment_dir = tempfile::TempDir::new().unwrap();
+        let prefix_record = download_and_get_prefix_record(
+            environment_dir.path(),
+            "https://conda.anaconda.org/conda-forge/win-64/ruff-0.0.171-py310h298983d_0.conda"
+                .parse()
+                .unwrap(),
+            "25c755b97189ee066576b4ae3999d5e7ff4406d236b984742194e63941838dcd",
+        )
+        .await;
+        let name = prefix_record.repodata_record.package_record.name.clone();
+
+        let transaction = Transaction::from_current_and_desired(
+            vec![prefix_record.clone()],
+            vec![prefix_record.clone()],
+            Some(HashSet::from_iter(vec![name])),
+            Platform::current(),
+        )
+        .unwrap();
+
+        assert_matches!(
+            transaction.operations[0],
+            TransactionOperation::Change { .. }
+        );
+    }
+}
diff --git a/crates/rattler/src/install/unlink.rs b/crates/rattler/src/install/unlink.rs
@@ -227,59 +227,17 @@ mod tests {
         fs::{self, File},
         io::Write,
         path::Path,
-        str::FromStr,
     };
 
-    use rattler_conda_types::{Platform, PrefixRecord, RepoDataRecord, Version};
-    use url::Url;
+    use rattler_conda_types::{Platform, RepoDataRecord};
 
-    use crate::{
-        get_repodata_record,
-        install::{
-            empty_trash, link_package, unlink_package, InstallDriver, InstallOptions, PythonInfo,
-            Transaction,
-        },
-    };
-
-    async fn link_ruff(target_prefix: &Path, package_url: Url, sha256_hash: &str) -> PrefixRecord {
-        let package_path = tools::download_and_cache_file_async(package_url, sha256_hash)
-            .await
-            .unwrap();
-
-        let package_dir = tempfile::TempDir::new().unwrap();
-
-        // Create package cache
-        rattler_package_streaming::fs::extract(&package_path, package_dir.path()).unwrap();
-
-        let py_info =
-            PythonInfo::from_version(&Version::from_str("3.10").unwrap(), None, Platform::Linux64)
-                .unwrap();
-        let install_options = InstallOptions {
-            python_info: Some(py_info),
-            ..InstallOptions::default()
-        };
-
-        let install_driver = InstallDriver::default();
-        // Link the package
-        let paths = link_package(
-            package_dir.path(),
-            target_prefix,
-            &install_driver,
-            install_options,
-        )
-        .await
-        .unwrap();
-
-        let repodata_record = get_repodata_record(&package_path);
-        // Construct a PrefixRecord for the package
-
-        PrefixRecord::from_repodata_record(repodata_record, None, None, paths, None, None)
-    }
+    use crate::install::test_utils::download_and_get_prefix_record;
+    use crate::install::{empty_trash, unlink_package, InstallDriver, Transaction};
 
     #[tokio::test]
     async fn test_unlink_package() {
         let environment_dir = tempfile::TempDir::new().unwrap();
-        let prefix_record = link_ruff(
+        let prefix_record = download_and_get_prefix_record(
             environment_dir.path(),
             "https://conda.anaconda.org/conda-forge/win-64/ruff-0.0.171-py310h298983d_0.conda"
                 .parse()
@@ -308,6 +266,7 @@ mod tests {
         let transaction = Transaction::from_current_and_desired(
             vec![prefix_record.clone()],
             Vec::<RepoDataRecord>::new().into_iter(),
+            None,
             Platform::current(),
         )
         .unwrap();
@@ -328,7 +287,7 @@ mod tests {
     #[tokio::test]
     async fn test_unlink_package_python_noarch() {
         let target_prefix = tempfile::TempDir::new().unwrap();
-        let prefix_record = link_ruff(
+        let prefix_record = download_and_get_prefix_record(
             target_prefix.path(),
             "https://conda.anaconda.org/conda-forge/noarch/pytweening-1.0.4-pyhd8ed1ab_0.tar.bz2"
                 .parse()
@@ -370,6 +329,7 @@ mod tests {
         let transaction = Transaction::from_current_and_desired(
             vec![prefix_record.clone()],
             Vec::<RepoDataRecord>::new().into_iter(),
+            None,
             Platform::current(),
         )
         .unwrap();
@@ -404,6 +364,10 @@ mod tests {
     #[cfg(windows)]
     #[tokio::test]
     async fn test_unlink_package_in_use() {
+        use crate::get_repodata_record;
+        use crate::install::link_package;
+        use crate::install::InstallOptions;
+        use rattler_conda_types::PrefixRecord;
         use std::{
             env::{join_paths, split_paths, var_os},
             io::{BufRead, BufReader},
