Disable concurrent progress bars in Jupyter Notebooks (#3890)

## Summary

Resolves https://github.com/astral-sh/uv/issues/3887 by disabling the
new progress output when the `JPY_SESSION_NAME` environment variable is
detected.
This commit is contained in:
Ibraheem Ahmed 2024-05-28 17:05:11 -04:00 committed by GitHub
parent 47db418ba2
commit 038af6e658
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 98 additions and 36 deletions

View File

@ -1,3 +1,4 @@
use std::env;
use std::sync::{Arc, Mutex};
use std::time::Duration;
@ -17,8 +18,18 @@ use crate::printer::Printer;
struct ProgressReporter {
printer: Printer,
root: ProgressBar,
mode: ProgressMode,
}
#[derive(Debug)]
enum ProgressMode {
// Reports top-level progress.
Single,
// Reports progress of all concurrent download/build/checkout processes.
Multi {
multi_progress: MultiProgress,
state: Arc<Mutex<BarState>>,
},
}
#[derive(Default, Debug)]
@ -42,11 +53,41 @@ impl BarState {
}
impl ProgressReporter {
fn new(mut root: ProgressBar, printer: Printer) -> ProgressReporter {
let mode = if env::var("JPY_SESSION_NAME").is_ok() {
// Disable concurrent progress bars when running inside a Jupyter notebook
// because the Jupyter terminal does not support clearing previous lines.
// See: https://github.com/astral-sh/uv/issues/3887.
ProgressMode::Single
} else {
let multi_progress = MultiProgress::with_draw_target(printer.target());
root = multi_progress.add(root);
ProgressMode::Multi {
state: Arc::default(),
multi_progress,
}
};
ProgressReporter {
printer,
root,
mode,
}
}
fn on_build_start(&self, source: &BuildableSource) -> usize {
let mut state = self.state.lock().unwrap();
let ProgressMode::Multi {
multi_progress,
state,
} = &self.mode
else {
return 0;
};
let mut state = state.lock().unwrap();
let id = state.id();
let progress = self.multi_progress.insert_before(
let progress = multi_progress.insert_before(
&self.root,
ProgressBar::with_draw_target(None, self.printer.target()),
);
@ -64,8 +105,12 @@ impl ProgressReporter {
}
fn on_build_complete(&self, source: &BuildableSource, id: usize) {
let ProgressMode::Multi { state, .. } = &self.mode else {
return;
};
let progress = {
let mut state = self.state.lock().unwrap();
let mut state = state.lock().unwrap();
state.headers -= 1;
state.bars.remove(&id).unwrap()
};
@ -78,13 +123,21 @@ impl ProgressReporter {
}
fn on_download_start(&self, name: &PackageName, size: Option<u64>) -> usize {
let mut state = self.state.lock().unwrap();
let ProgressMode::Multi {
multi_progress,
state,
} = &self.mode
else {
return 0;
};
let mut state = state.lock().unwrap();
// Preserve ascending order.
let position = size.map_or(0, |size| state.sizes.partition_point(|&len| len < size));
state.sizes.insert(position, size.unwrap_or(0));
let progress = self.multi_progress.insert(
let progress = multi_progress.insert(
// Make sure not to reorder the initial "Downloading..." bar, or any previous bars.
position + 1 + state.headers,
ProgressBar::with_draw_target(size, self.printer.target()),
@ -111,19 +164,35 @@ impl ProgressReporter {
}
fn on_download_progress(&self, id: usize, bytes: u64) {
self.state.lock().unwrap().bars[&id].inc(bytes);
let ProgressMode::Multi { state, .. } = &self.mode else {
return;
};
state.lock().unwrap().bars[&id].inc(bytes);
}
fn on_download_complete(&self, id: usize) {
let progress = self.state.lock().unwrap().bars.remove(&id).unwrap();
let ProgressMode::Multi { state, .. } = &self.mode else {
return;
};
let progress = state.lock().unwrap().bars.remove(&id).unwrap();
progress.finish_and_clear();
}
fn on_checkout_start(&self, url: &Url, rev: &str) -> usize {
let mut state = self.state.lock().unwrap();
let ProgressMode::Multi {
multi_progress,
state,
} = &self.mode
else {
return 0;
};
let mut state = state.lock().unwrap();
let id = state.id();
let progress = self.multi_progress.insert_before(
let progress = multi_progress.insert_before(
&self.root,
ProgressBar::with_draw_target(None, self.printer.target()),
);
@ -143,8 +212,12 @@ impl ProgressReporter {
}
fn on_checkout_complete(&self, url: &Url, rev: &str, id: usize) {
let ProgressMode::Multi { state, .. } = &self.mode else {
return;
};
let progress = {
let mut state = self.state.lock().unwrap();
let mut state = state.lock().unwrap();
state.headers -= 1;
state.bars.remove(&id).unwrap()
};
@ -165,24 +238,16 @@ pub(crate) struct DownloadReporter {
impl From<Printer> for DownloadReporter {
fn from(printer: Printer) -> Self {
let multi_progress = MultiProgress::with_draw_target(printer.target());
let progress = multi_progress.add(ProgressBar::with_draw_target(None, printer.target()));
progress.enable_steady_tick(Duration::from_millis(200));
progress.set_style(
let root = ProgressBar::with_draw_target(None, printer.target());
root.enable_steady_tick(Duration::from_millis(200));
root.set_style(
ProgressStyle::with_template("{spinner:.white} {msg:.dim} ({pos}/{len})")
.unwrap()
.tick_strings(&["", "", "", "", "", "", "", "", "", ""]),
);
progress.set_message("Downloading packages...");
let reporter = ProgressReporter {
printer,
multi_progress,
root: progress,
state: Arc::default(),
};
root.set_message("Downloading packages...");
let reporter = ProgressReporter::new(root, printer);
Self { reporter }
}
}
@ -201,6 +266,9 @@ impl uv_installer::DownloadReporter for DownloadReporter {
}
fn on_complete(&self) {
// Need an extra call to `set_message` here to fully clear avoid leaving ghost output
// in Jupyter notebooks.
self.reporter.root.set_message("");
self.reporter.root.finish_and_clear();
}
@ -248,9 +316,7 @@ impl ResolverReporter {
impl From<Printer> for ResolverReporter {
fn from(printer: Printer) -> Self {
let multi_progress = MultiProgress::with_draw_target(printer.target());
let root = multi_progress.add(ProgressBar::with_draw_target(None, printer.target()));
let root = ProgressBar::with_draw_target(None, printer.target());
root.enable_steady_tick(Duration::from_millis(200));
root.set_style(
ProgressStyle::with_template("{spinner:.white} {wide_msg:.dim}")
@ -259,14 +325,8 @@ impl From<Printer> for ResolverReporter {
);
root.set_message("Resolving dependencies...");
let reporter = ProgressReporter {
root,
printer,
multi_progress,
state: Arc::default(),
};
ResolverReporter { reporter }
let reporter = ProgressReporter::new(root, printer);
Self { reporter }
}
}
@ -283,6 +343,7 @@ impl uv_resolver::ResolverReporter for ResolverReporter {
}
fn on_complete(&self) {
self.reporter.root.set_message("");
self.reporter.root.finish_and_clear();
}
@ -376,6 +437,7 @@ impl uv_installer::InstallReporter for InstallReporter {
}
fn on_install_complete(&self) {
self.progress.set_message("");
self.progress.finish_and_clear();
}
}