Add progress bar for uv cache clean (#8857)

## Summary

Closes #8786
This commit is contained in:
Jo 2024-11-07 00:43:43 +08:00 committed by GitHub
parent 987d778631
commit dcd24b7c42
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 171 additions and 127 deletions

View file

@ -320,8 +320,8 @@ impl Cache {
}
/// Clear the cache, removing all entries.
pub fn clear(&self) -> Result<Removal, io::Error> {
rm_rf(&self.root)
pub fn clear(&self, reporter: Option<&dyn CleanReporter>) -> Result<Removal, io::Error> {
rm_rf(&self.root, reporter)
}
/// Remove a package from the cache.
@ -379,7 +379,7 @@ impl Cache {
let path = fs_err::canonicalize(entry.path())?;
if !after.contains(&path) && before.contains(&path) {
debug!("Removing dangling cache entry: {}", path.display());
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
}
}
}
@ -409,13 +409,13 @@ impl Cache {
if CacheBucket::iter().all(|bucket| entry.file_name() != bucket.to_str()) {
let path = entry.path();
debug!("Removing dangling cache bucket: {}", path.display());
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
}
} else {
// If the file is not a marker file, remove it.
let path = entry.path();
debug!("Removing dangling cache bucket: {}", path.display());
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
}
}
@ -427,7 +427,7 @@ impl Cache {
let entry = entry?;
let path = fs_err::canonicalize(entry.path())?;
debug!("Removing dangling cache environment: {}", path.display());
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
}
}
Err(err) if err.kind() == io::ErrorKind::NotFound => (),
@ -444,7 +444,7 @@ impl Cache {
let path = fs_err::canonicalize(entry.path())?;
if path.is_dir() {
debug!("Removing unzipped wheel entry: {}", path.display());
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
}
}
}
@ -472,10 +472,10 @@ impl Cache {
if path.is_dir() {
debug!("Removing unzipped built wheel entry: {}", path.display());
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
} else if path.is_symlink() {
debug!("Removing unzipped built wheel entry: {}", path.display());
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
}
}
}
@ -505,7 +505,7 @@ impl Cache {
let path = fs_err::canonicalize(entry.path())?;
if !references.contains(&path) {
debug!("Removing dangling cache archive: {}", path.display());
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
}
}
}
@ -517,6 +517,15 @@ impl Cache {
}
}
pub trait CleanReporter: Send + Sync {
/// Called after one file or directory is removed.
fn on_clean(&self);
/// Called after a package is cleaned.
fn on_clean_package(&self, _package: &str, _removal: &Removal) {}
/// Called after all files and directories are removed.
fn on_complete(&self);
}
/// The different kinds of data in the cache are stored in different bucket, which in our case
/// are subdirectories of the cache root.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
@ -800,32 +809,32 @@ impl CacheBucket {
Self::Wheels => {
// For `pypi` wheels, we expect a directory per package (indexed by name).
let root = cache.bucket(self).join(WheelCacheKind::Pypi);
summary += rm_rf(root.join(name.to_string()))?;
summary += rm_rf(root.join(name.to_string()), None)?;
// For alternate indices, we expect a directory for every index (under an `index`
// subdirectory), followed by a directory per package (indexed by name).
let root = cache.bucket(self).join(WheelCacheKind::Index);
for directory in directories(root) {
summary += rm_rf(directory.join(name.to_string()))?;
summary += rm_rf(directory.join(name.to_string()), None)?;
}
// For direct URLs, we expect a directory for every URL, followed by a
// directory per package (indexed by name).
let root = cache.bucket(self).join(WheelCacheKind::Url);
for directory in directories(root) {
summary += rm_rf(directory.join(name.to_string()))?;
summary += rm_rf(directory.join(name.to_string()), None)?;
}
}
Self::SourceDistributions => {
// For `pypi` wheels, we expect a directory per package (indexed by name).
let root = cache.bucket(self).join(WheelCacheKind::Pypi);
summary += rm_rf(root.join(name.to_string()))?;
summary += rm_rf(root.join(name.to_string()), None)?;
// For alternate indices, we expect a directory for every index (under an `index`
// subdirectory), followed by a directory per package (indexed by name).
let root = cache.bucket(self).join(WheelCacheKind::Index);
for directory in directories(root) {
summary += rm_rf(directory.join(name.to_string()))?;
summary += rm_rf(directory.join(name.to_string()), None)?;
}
// For direct URLs, we expect a directory for every URL, followed by a
@ -834,7 +843,7 @@ impl CacheBucket {
let root = cache.bucket(self).join(WheelCacheKind::Url);
for url in directories(root) {
if directories(&url).any(|version| is_match(&version, name)) {
summary += rm_rf(url)?;
summary += rm_rf(url, None)?;
}
}
@ -844,7 +853,7 @@ impl CacheBucket {
let root = cache.bucket(self).join(WheelCacheKind::Path);
for path in directories(root) {
if directories(&path).any(|version| is_match(&version, name)) {
summary += rm_rf(path)?;
summary += rm_rf(path, None)?;
}
}
@ -855,7 +864,7 @@ impl CacheBucket {
for repository in directories(root) {
for sha in directories(repository) {
if is_match(&sha, name) {
summary += rm_rf(sha)?;
summary += rm_rf(sha, None)?;
}
}
}
@ -863,20 +872,20 @@ impl CacheBucket {
Self::Simple => {
// For `pypi` wheels, we expect a rkyv file per package, indexed by name.
let root = cache.bucket(self).join(WheelCacheKind::Pypi);
summary += rm_rf(root.join(format!("{name}.rkyv")))?;
summary += rm_rf(root.join(format!("{name}.rkyv")), None)?;
// For alternate indices, we expect a directory for every index (under an `index`
// subdirectory), followed by a directory per package (indexed by name).
let root = cache.bucket(self).join(WheelCacheKind::Index);
for directory in directories(root) {
summary += rm_rf(directory.join(format!("{name}.rkyv")))?;
summary += rm_rf(directory.join(format!("{name}.rkyv")), None)?;
}
}
Self::FlatIndex => {
// We can't know if the flat index includes a package, so we just remove the entire
// cache entry.
let root = cache.bucket(self);
summary += rm_rf(root)?;
summary += rm_rf(root, None)?;
}
Self::Git => {
// Nothing to do.

View file

@ -5,11 +5,13 @@
use std::io;
use std::path::Path;
use crate::CleanReporter;
/// Remove a file or directory and all its contents, returning a [`Removal`] with
/// the number of files and directories removed, along with a total byte count.
pub fn rm_rf(path: impl AsRef<Path>) -> io::Result<Removal> {
pub fn rm_rf(path: impl AsRef<Path>, reporter: Option<&dyn CleanReporter>) -> io::Result<Removal> {
let mut removal = Removal::default();
removal.rm_rf(path.as_ref())?;
removal.rm_rf(path.as_ref(), reporter)?;
Ok(removal)
}
@ -28,7 +30,7 @@ pub struct Removal {
impl Removal {
/// Recursively remove a file or directory and all its contents.
fn rm_rf(&mut self, path: &Path) -> io::Result<()> {
fn rm_rf(&mut self, path: &Path, reporter: Option<&dyn CleanReporter>) -> io::Result<()> {
let metadata = match fs_err::symlink_metadata(path) {
Ok(metadata) => metadata,
Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(()),
@ -47,6 +49,8 @@ impl Removal {
remove_file(path)?;
}
reporter.map(CleanReporter::on_clean);
return Ok(());
}
@ -61,7 +65,7 @@ impl Removal {
if set_readable(dir).unwrap_or(false) {
// Retry the operation; if we _just_ `self.rm_rf(dir)` and continue,
// `walkdir` may give us duplicate entries for the directory.
return self.rm_rf(path);
return self.rm_rf(path, reporter);
}
}
}
@ -88,8 +92,12 @@ impl Removal {
}
remove_file(entry.path())?;
}
reporter.map(CleanReporter::on_clean);
}
reporter.map(CleanReporter::on_complete);
Ok(())
}
}

View file

@ -1995,8 +1995,8 @@ pub fn prune(cache: &Cache) -> Result<Removal, Error> {
"Removing dangling source revision: {}",
sibling.path().display()
);
removal +=
uv_cache::rm_rf(sibling.path()).map_err(Error::CacheWrite)?;
removal += uv_cache::rm_rf(sibling.path(), None)
.map_err(Error::CacheWrite)?;
}
}
}
@ -2020,8 +2020,8 @@ pub fn prune(cache: &Cache) -> Result<Removal, Error> {
"Removing dangling source revision: {}",
sibling.path().display()
);
removal +=
uv_cache::rm_rf(sibling.path()).map_err(Error::CacheWrite)?;
removal += uv_cache::rm_rf(sibling.path(), None)
.map_err(Error::CacheWrite)?;
}
}
}

View file

@ -93,6 +93,7 @@ tracing-subscriber = { workspace = true, features = ["json"] }
tracing-tree = { workspace = true }
unicode-width = { workspace = true }
url = { workspace = true }
walkdir = { workspace = true }
which = { workspace = true }
zip = { workspace = true }

View file

@ -3,10 +3,11 @@ use std::fmt::Write;
use anyhow::{Context, Result};
use owo_colors::OwoColorize;
use uv_cache::Cache;
use uv_cache::{Cache, CleanReporter, Removal};
use uv_fs::Simplified;
use uv_normalize::PackageName;
use crate::commands::reporters::{CleaningDirectoryReporter, CleaningPackageReporter};
use crate::commands::{human_readable_bytes, ExitStatus};
use crate::printer::Printer;
@ -25,101 +26,65 @@ pub(crate) fn cache_clean(
return Ok(ExitStatus::Success);
}
if packages.is_empty() {
let summary = if packages.is_empty() {
writeln!(
printer.stderr(),
"Clearing cache at: {}",
cache.root().user_display().cyan()
)?;
let summary = cache.clear().with_context(|| {
format!("Failed to clear cache at: {}", cache.root().user_display())
})?;
let num_paths = walkdir::WalkDir::new(cache.root()).into_iter().count();
let reporter = CleaningDirectoryReporter::new(printer, num_paths);
// Write a summary of the number of files and directories removed.
match (summary.num_files, summary.num_dirs) {
(0, 0) => {
write!(printer.stderr(), "No cache entries found")?;
}
(0, 1) => {
write!(printer.stderr(), "Removed 1 directory")?;
}
(0, num_dirs_removed) => {
write!(printer.stderr(), "Removed {num_dirs_removed} directories")?;
}
(1, _) => {
write!(printer.stderr(), "Removed 1 file")?;
}
(num_files_removed, _) => {
write!(printer.stderr(), "Removed {num_files_removed} files")?;
}
}
// If any, write a summary of the total byte count removed.
if summary.total_bytes > 0 {
let bytes = if summary.total_bytes < 1024 {
format!("{}B", summary.total_bytes)
} else {
let (bytes, unit) = human_readable_bytes(summary.total_bytes);
format!("{bytes:.1}{unit}")
};
write!(printer.stderr(), " ({})", bytes.green())?;
}
writeln!(printer.stderr())?;
cache
.clear(Some(&reporter))
.with_context(|| format!("Failed to clear cache at: {}", cache.root().user_display()))?
} else {
let reporter = CleaningPackageReporter::new(printer, packages.len());
let mut summary = Removal::default();
for package in packages {
let summary = cache.remove(package)?;
let removed = cache.remove(package)?;
summary += removed;
// Write a summary of the number of files and directories removed.
match (summary.num_files, summary.num_dirs) {
(0, 0) => {
write!(
printer.stderr(),
"No cache entries found for {}",
package.cyan()
)?;
}
(0, 1) => {
write!(
printer.stderr(),
"Removed 1 directory for {}",
package.cyan()
)?;
}
(0, num_dirs_removed) => {
write!(
printer.stderr(),
"Removed {num_dirs_removed} directories for {}",
package.cyan()
)?;
}
(1, _) => {
write!(printer.stderr(), "Removed 1 file for {}", package.cyan())?;
}
(num_files_removed, _) => {
write!(
printer.stderr(),
"Removed {num_files_removed} files for {}",
package.cyan()
)?;
}
}
reporter.on_clean_package(package.as_str(), &summary);
}
reporter.on_complete();
// If any, write a summary of the total byte count removed.
if summary.total_bytes > 0 {
let bytes = if summary.total_bytes < 1024 {
format!("{}B", summary.total_bytes)
} else {
let (bytes, unit) = human_readable_bytes(summary.total_bytes);
format!("{bytes:.1}{unit}")
};
write!(printer.stderr(), " ({})", bytes.green())?;
}
summary
};
writeln!(printer.stderr())?;
// Write a summary of the number of files and directories removed.
match (summary.num_files, summary.num_dirs) {
(0, 0) => {
write!(printer.stderr(), "No cache entries found")?;
}
(0, 1) => {
write!(printer.stderr(), "Removed 1 directory")?;
}
(0, num_dirs_removed) => {
write!(printer.stderr(), "Removed {num_dirs_removed} directories")?;
}
(1, _) => {
write!(printer.stderr(), "Removed 1 file")?;
}
(num_files_removed, _) => {
write!(printer.stderr(), "Removed {num_files_removed} files")?;
}
}
// If any, write a summary of the total byte count removed.
if summary.total_bytes > 0 {
let bytes = if summary.total_bytes < 1024 {
format!("{}B", summary.total_bytes)
} else {
let (bytes, unit) = human_readable_bytes(summary.total_bytes);
format!("{bytes:.1}{unit}")
};
write!(printer.stderr(), " ({})", bytes.green())?;
}
writeln!(printer.stderr())?;
Ok(ExitStatus::Success)
}

View file

@ -6,7 +6,7 @@ use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use owo_colors::OwoColorize;
use rustc_hash::FxHashMap;
use url::Url;
use uv_cache::Removal;
use uv_distribution_types::{
BuildableSource, CachedDist, DistributionMetadata, Name, SourceDist, VersionOrUrlRef,
};
@ -528,6 +528,69 @@ impl uv_publish::Reporter for PublishReporter {
}
}
#[derive(Debug)]
pub(crate) struct CleaningDirectoryReporter {
bar: ProgressBar,
}
impl CleaningDirectoryReporter {
/// Initialize a [`CleaningDirectoryReporter`] for cleaning the cache directory.
pub(crate) fn new(printer: Printer, max: usize) -> Self {
let bar = ProgressBar::with_draw_target(Some(max as u64), printer.target());
bar.set_style(
ProgressStyle::with_template("{prefix} [{bar:20}] {percent}%")
.unwrap()
.progress_chars("=> "),
);
bar.set_prefix(format!("{}", "Cleaning".bold().cyan()));
Self { bar }
}
}
impl uv_cache::CleanReporter for CleaningDirectoryReporter {
fn on_clean(&self) {
self.bar.inc(1);
}
fn on_complete(&self) {
self.bar.finish_and_clear();
}
}
pub(crate) struct CleaningPackageReporter {
bar: ProgressBar,
}
impl CleaningPackageReporter {
/// Initialize a [`CleaningPackageReporter`] for cleaning packages from the cache.
pub(crate) fn new(printer: Printer, max: usize) -> Self {
let bar = ProgressBar::with_draw_target(Some(max as u64), printer.target());
bar.set_style(
ProgressStyle::with_template("{prefix} [{bar:20}] {pos}/{len}{msg}")
.unwrap()
.progress_chars("=> "),
);
bar.set_prefix(format!("{}", "Cleaning".bold().cyan()));
Self { bar }
}
}
impl uv_cache::CleanReporter for CleaningPackageReporter {
fn on_clean(&self) {}
fn on_clean_package(&self, package: &str, removal: &Removal) {
self.bar.inc(1);
self.bar.set_message(format!(
": {}, {} files {} folders removed",
package, removal.num_files, removal.num_dirs,
));
}
fn on_complete(&self) {
self.bar.finish_and_clear();
}
}
/// Like [`std::fmt::Display`], but with colors.
trait ColorDisplay {
fn to_color_string(&self) -> String;

View file

@ -79,7 +79,7 @@ fn clean_package_pypi() -> Result<()> {
----- stderr -----
DEBUG uv [VERSION] ([COMMIT] DATE)
DEBUG Removing dangling cache entry: [CACHE_DIR]/archive-v0/[ENTRY]
Removed 12 files for iniconfig ([SIZE])
Removed 12 files ([SIZE])
"###);
// Assert that the `.rkyv` file is removed for `iniconfig`.
@ -152,7 +152,7 @@ fn clean_package_index() -> Result<()> {
----- stderr -----
DEBUG uv [VERSION] ([COMMIT] DATE)
DEBUG Removing dangling cache entry: [CACHE_DIR]/archive-v0/[ENTRY]
Removed 12 files for iniconfig ([SIZE])
Removed 12 files ([SIZE])
"###);
// Assert that the `.rkyv` file is removed for `iniconfig`.

View file

@ -1410,7 +1410,7 @@ fn install_url_source_dist_cached() -> Result<()> {
----- stdout -----
----- stderr -----
Removed 19 files for source-distribution ([SIZE])
Removed 19 files ([SIZE])
"###
);
@ -1505,7 +1505,7 @@ fn install_git_source_dist_cached() -> Result<()> {
----- stdout -----
----- stderr -----
No cache entries found for werkzeug
No cache entries found
"###
);
@ -1605,7 +1605,7 @@ fn install_registry_source_dist_cached() -> Result<()> {
----- stdout -----
----- stderr -----
Removed 20 files for source-distribution ([SIZE])
Removed 20 files ([SIZE])
"###
);
@ -1702,7 +1702,7 @@ fn install_path_source_dist_cached() -> Result<()> {
----- stdout -----
----- stderr -----
Removed 19 files for source-distribution ([SIZE])
Removed 19 files ([SIZE])
"###
);
@ -1787,13 +1787,10 @@ fn install_path_built_dist_cached() -> Result<()> {
let filters = if cfg!(windows) {
// We do not display sizes on Windows
[(
"Removed 1 file for tomli",
"Removed 1 file for tomli ([SIZE])",
)]
.into_iter()
.chain(context.filters())
.collect()
[("Removed 1 file", "Removed 1 file ([SIZE])")]
.into_iter()
.chain(context.filters())
.collect()
} else {
context.filters()
};
@ -1804,7 +1801,7 @@ fn install_path_built_dist_cached() -> Result<()> {
----- stdout -----
----- stderr -----
Removed 11 files for tomli ([SIZE])
Removed 11 files ([SIZE])
"###
);
@ -1892,7 +1889,7 @@ fn install_url_built_dist_cached() -> Result<()> {
----- stdout -----
----- stderr -----
Removed 43 files for tqdm ([SIZE])
Removed 43 files ([SIZE])
"###
);