Process completed Python installs and uninstalls as a stream (#5203)

## Summary This ensures that we process Python installs and uninstalls as soon as they complete, rather than waiting for them all to complete, then processing them sequentially. In practice, it shouldn't be much of a difference (since the processing is code is fairly light), but it strikes me as more correct.
astral-sh · Jul 19, 2024 · fb1a529 · fb1a529
1 parent ef56df2
commit fb1a529
Show file tree

Hide file tree

Showing 5 changed files with 138 additions and 129 deletions.
diff --git a/crates/pep508-rs/src/marker.rs b/crates/pep508-rs/src/marker.rs
@@ -1631,7 +1631,7 @@ impl Display for MarkerExpression {
                 {
                     return write!(f, "{key} {op} '{version}.*'");
                 }
-                write!(f, "{key} {op} '{version}'",)
+                write!(f, "{key} {op} '{version}'")
             }
             MarkerExpression::VersionInverted {
                 version,

diff --git a/crates/platform-tags/src/tags.rs b/crates/platform-tags/src/tags.rs
@@ -327,7 +327,7 @@ impl Implementation {
                 } else if gil_disabled {
                     // https://peps.python.org/pep-0703/#build-configuration-changes
                     // Python 3.13+ only, but it makes more sense to just rely on the sysconfig var.
-                    format!("cp{}{}t", python_version.0, python_version.1,)
+                    format!("cp{}{}t", python_version.0, python_version.1)
                 } else {
                     format!(
                         "cp{}{}{}",

diff --git a/crates/uv-resolver/src/resolution/requirements_txt.rs b/crates/uv-resolver/src/resolution/requirements_txt.rs
@@ -99,7 +99,7 @@ impl RequirementsTxtDist {
 
         if self.extras.is_empty() || !include_extras {
             if let Some(markers) = self.markers.as_ref().filter(|_| include_markers) {
-                Cow::Owned(format!("{} ; {}", self.dist.verbatim(), markers,))
+                Cow::Owned(format!("{} ; {}", self.dist.verbatim(), markers))
             } else {
                 self.dist.verbatim()
             }

diff --git a/crates/uv/src/commands/python/install.rs b/crates/uv/src/commands/python/install.rs
@@ -1,12 +1,12 @@
-use std::collections::BTreeSet;
-use std::fmt::Write;
-use std::path::PathBuf;
-
 use anyhow::Result;
 use fs_err as fs;
+use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use itertools::Itertools;
 use owo_colors::OwoColorize;
+use std::collections::BTreeSet;
+use std::fmt::Write;
+use std::path::PathBuf;
 use tracing::debug;
 use uv_cache::Cache;
 use uv_client::Connectivity;
@@ -144,20 +144,21 @@ pub(crate) async fn install(
 
     let reporter = PythonDownloadReporter::new(printer, downloads.len() as u64);
 
-    let results = futures::stream::iter(downloads.iter())
-        .map(|download| async {
-            let result = download
-                .fetch(&client, installations_dir, Some(&reporter))
-                .await;
-            (download.key(), result)
-        })
-        .buffered(4)
-        .collect::<Vec<_>>()
-        .await;
+    let mut tasks = FuturesUnordered::new();
+    for download in &downloads {
+        tasks.push(async {
+            (
+                download.key(),
+                download
+                    .fetch(&client, installations_dir, Some(&reporter))
+                    .await,
+            )
+        });
+    }
 
     let mut installed = vec![];
-    let mut failed = false;
-    for (key, result) in results {
+    let mut errors = vec![];
+    while let Some((key, result)) = tasks.next().await {
         match result {
             Ok(download) => {
                 let path = match download {
@@ -173,66 +174,72 @@ pub(crate) async fn install(
                 managed.ensure_externally_managed()?;
             }
             Err(err) => {
-                failed = true;
-                writeln!(printer.stderr(), "Failed to install {}: {err}", key.green())?;
+                errors.push((key, err));
             }
         }
     }
 
-    if failed {
-        if downloads.len() > 1 {
-            writeln!(printer.stderr(), "Failed to install some Python versions")?;
+    if !installed.is_empty() {
+        if let [installed] = installed.as_slice() {
+            // Ex) "Installed Python 3.9.7 in 1.68s"
+            writeln!(
+                printer.stderr(),
+                "{}",
+                format!(
+                    "Installed {} {}",
+                    format!("Python {}", installed.version()).bold(),
+                    format!("in {}", elapsed(start.elapsed())).dimmed()
+                )
+                .dimmed()
+            )?;
+        } else {
+            // Ex) "Installed 2 versions in 1.68s"
+            let s = if installed.len() == 1 { "" } else { "s" };
+            writeln!(
+                printer.stderr(),
+                "{}",
+                format!(
+                    "Installed {} {}",
+                    format!("{} version{s}", installed.len()).bold(),
+                    format!("in {}", elapsed(start.elapsed())).dimmed()
+                )
+                .dimmed()
+            )?;
         }
-        return Ok(ExitStatus::Failure);
-    }
 
-    if let [installed] = installed.as_slice() {
-        // Ex) "Installed Python 3.9.7 in 1.68s"
-        writeln!(
-            printer.stderr(),
-            "{}",
-            format!(
-                "Installed {} {}",
-                format!("Python {}", installed.version()).bold(),
-                format!("in {}", elapsed(start.elapsed())).dimmed()
-            )
-            .dimmed()
-        )?;
-    } else {
-        // Ex) "Installed 2 versions in 1.68s"
-        let s = if installed.len() == 1 { "" } else { "s" };
-        writeln!(
-            printer.stderr(),
-            "{}",
-            format!(
-                "Installed {} {}",
-                format!("{} version{s}", installed.len()).bold(),
-                format!("in {}", elapsed(start.elapsed())).dimmed()
-            )
-            .dimmed()
-        )?;
+        for event in uninstalled
+            .into_iter()
+            .map(|key| ChangeEvent {
+                key,
+                kind: ChangeEventKind::Removed,
+            })
+            .chain(installed.into_iter().map(|key| ChangeEvent {
+                key,
+                kind: ChangeEventKind::Added,
+            }))
+            .sorted_unstable_by(|a, b| a.key.cmp(&b.key).then_with(|| a.kind.cmp(&b.kind)))
+        {
+            match event.kind {
+                ChangeEventKind::Added => {
+                    writeln!(printer.stderr(), " {} {}", "+".green(), event.key.bold())?;
+                }
+                ChangeEventKind::Removed => {
+                    writeln!(printer.stderr(), " {} {}", "-".red(), event.key.bold())?;
+                }
+            }
+        }
     }
 
-    for event in uninstalled
-        .into_iter()
-        .map(|key| ChangeEvent {
-            key,
-            kind: ChangeEventKind::Removed,
-        })
-        .chain(installed.into_iter().map(|key| ChangeEvent {
-            key,
-            kind: ChangeEventKind::Added,
-        }))
-        .sorted_unstable_by(|a, b| a.key.cmp(&b.key).then_with(|| a.kind.cmp(&b.kind)))
-    {
-        match event.kind {
-            ChangeEventKind::Added => {
-                writeln!(printer.stderr(), " {} {}", "+".green(), event.key.bold(),)?;
-            }
-            ChangeEventKind::Removed => {
-                writeln!(printer.stderr(), " {} {}", "-".red(), event.key.bold(),)?;
-            }
+    if !errors.is_empty() {
+        for (key, err) in errors {
+            writeln!(
+                printer.stderr(),
+                "Failed to install {}: {}",
+                key.green(),
+                err
+            )?;
         }
+        return Ok(ExitStatus::Failure);
     }
 
     Ok(ExitStatus::Success)

diff --git a/crates/uv/src/commands/python/uninstall.rs b/crates/uv/src/commands/python/uninstall.rs
@@ -2,6 +2,7 @@ use std::collections::BTreeSet;
 use std::fmt::Write;
 
 use anyhow::Result;
+use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use itertools::Itertools;
 use owo_colors::OwoColorize;
@@ -93,82 +94,83 @@ pub(crate) async fn uninstall(
         return Ok(ExitStatus::Failure);
     }
 
-    let results = futures::stream::iter(matching_installations.iter())
-        .map(|installation| async {
+    let mut tasks = FuturesUnordered::new();
+    for installation in &matching_installations {
+        tasks.push(async {
             (
                 installation.key(),
                 fs_err::tokio::remove_dir_all(installation.path()).await,
             )
-        })
-        .buffered(4)
-        .collect::<Vec<_>>()
-        .await;
+        });
+    }
 
     let mut uninstalled = vec![];
-    let mut failed = false;
-    for (key, result) in results {
+    let mut errors = vec![];
+    while let Some((key, result)) = tasks.next().await {
         if let Err(err) = result {
-            failed = true;
-            writeln!(
-                printer.stderr(),
-                "Failed to uninstall {}: {err}",
-                key.green()
-            )?;
+            errors.push((key.clone(), err));
         } else {
             uninstalled.push(key.clone());
         }
     }
 
-    if failed {
-        if matching_installations.len() > 1 {
-            writeln!(printer.stderr(), "Failed to uninstall some Python versions")?;
+    if !uninstalled.is_empty() {
+        if let [uninstalled] = uninstalled.as_slice() {
+            // Ex) "Uninstalled Python 3.9.7 in 1.68s"
+            writeln!(
+                printer.stderr(),
+                "{}",
+                format!(
+                    "Uninstalled {} {}",
+                    format!("Python {}", uninstalled.version()).bold(),
+                    format!("in {}", elapsed(start.elapsed())).dimmed()
+                )
+                .dimmed()
+            )?;
+        } else {
+            // Ex) "Uninstalled 2 versions in 1.68s"
+            let s = if uninstalled.len() == 1 { "" } else { "s" };
+            writeln!(
+                printer.stderr(),
+                "{}",
+                format!(
+                    "Uninstalled {} {}",
+                    format!("{} version{s}", uninstalled.len()).bold(),
+                    format!("in {}", elapsed(start.elapsed())).dimmed()
+                )
+                .dimmed()
+            )?;
         }
-        return Ok(ExitStatus::Failure);
-    }
 
-    if let [uninstalled] = uninstalled.as_slice() {
-        // Ex) "Uninstalled Python 3.9.7 in 1.68s"
-        writeln!(
-            printer.stderr(),
-            "{}",
-            format!(
-                "Uninstalled {} {}",
-                format!("Python {}", uninstalled.version()).bold(),
-                format!("in {}", elapsed(start.elapsed())).dimmed()
-            )
-            .dimmed()
-        )?;
-    } else {
-        // Ex) "Uninstalled 2 versions in 1.68s"
-        let s = if uninstalled.len() == 1 { "" } else { "s" };
-        writeln!(
-            printer.stderr(),
-            "{}",
-            format!(
-                "Uninstalled {} {}",
-                format!("{} version{s}", uninstalled.len()).bold(),
-                format!("in {}", elapsed(start.elapsed())).dimmed()
-            )
-            .dimmed()
-        )?;
+        for event in uninstalled
+            .into_iter()
+            .map(|key| ChangeEvent {
+                key,
+                kind: ChangeEventKind::Removed,
+            })
+            .sorted_unstable_by(|a, b| a.key.cmp(&b.key).then_with(|| a.kind.cmp(&b.kind)))
+        {
+            match event.kind {
+                ChangeEventKind::Added => {
+                    writeln!(printer.stderr(), " {} {}", "+".green(), event.key.bold())?;
+                }
+                ChangeEventKind::Removed => {
+                    writeln!(printer.stderr(), " {} {}", "-".red(), event.key.bold())?;
+                }
+            }
+        }
     }
 
-    for event in uninstalled
-        .into_iter()
-        .map(|key| ChangeEvent {
-            key,
-            kind: ChangeEventKind::Removed,
-        })
-        .sorted_unstable_by(|a, b| a.key.cmp(&b.key).then_with(|| a.kind.cmp(&b.kind)))
-    {
-        match event.kind {
-            ChangeEventKind::Added => {
-                writeln!(printer.stderr(), " {} {}", "+".green(), event.key.bold(),)?;
-            }
-            ChangeEventKind::Removed => {
-                writeln!(printer.stderr(), " {} {}", "-".red(), event.key.bold(),)?;
-            }
+    if !errors.is_empty() {
+        for (key, err) in errors {
+            writeln!(
+                printer.stderr(),
+                "Failed to uninstall {}: {}",
+                key.green(),
+                err
+            )?;
         }
+        return Ok(ExitStatus::Failure);
     }
 
     Ok(ExitStatus::Success)