#![warn(missing_docs)]
use std::{
collections::BTreeMap,
sync::{
atomic::{AtomicBool, AtomicUsize},
Arc, Mutex, Weak,
},
};
mod template;
mod ticker;
pub mod writer;
use template::{Template, TemplatePart};
use termsize::get_width;
use ticker::Ticker;
mod termsize;
const CLEAR_ANSI: &str = "\r\x1b[K";
const UP_ANSI: &str = "\x1b[F";
pub(crate) struct BarState {
len: u64,
pos: u64,
message: String,
template: Template,
created_at: std::time::Instant,
visible: bool,
need_redraw: bool,
}
fn duration_to_human(duration: std::time::Duration) -> String {
let elapsed = duration.as_secs();
let hours = elapsed / 3600;
let minutes = (elapsed % 3600) / 60;
let seconds = elapsed % 60;
format!("{}:{:02}:{:02}", hours, minutes, seconds)
}
fn bytes_to_human(bytes: u64) -> String {
const KB: u64 = 1024;
const MB: u64 = KB * 1024;
const GB: u64 = MB * 1024;
const TB: u64 = GB * 1024;
if bytes < KB {
format!("{} B", bytes)
} else if bytes < MB {
format!("{:.2} KiB", bytes as f64 / KB as f64)
} else if bytes < GB {
format!("{:.2} MiB", bytes as f64 / MB as f64)
} else if bytes < TB {
format!("{:.2} GiB", bytes as f64 / GB as f64)
} else {
format!("{:.2} TiB", bytes as f64 / TB as f64)
}
}
fn string_width(s: &str) -> usize {
#[cfg(feature = "unicode")]
{
unicode_width::UnicodeWidthStr::width(s)
}
#[cfg(not(feature = "unicode"))]
{
s.chars().count()
}
}
impl BarState {
pub fn render(&self) -> String {
let mut result = String::new();
let elapsed = std::time::Instant::now() - self.created_at;
let bytes_per_second = self.pos as f64 / elapsed.as_secs_f64();
for part in self.template.parts.iter() {
match part {
TemplatePart::Text(text) => {
result.push_str(text);
}
TemplatePart::Newline => {
result.push('\n');
}
TemplatePart::Message => {
result.push_str(&self.message);
}
TemplatePart::Elapsed => {
result.push_str(&duration_to_human(elapsed));
}
TemplatePart::Bytes => {
result.push_str(&bytes_to_human(self.pos));
}
TemplatePart::Pos => {
result.push_str(&self.pos.to_string());
}
TemplatePart::TotalBytes => {
result.push_str(&bytes_to_human(self.len));
}
TemplatePart::Total => {
result.push_str(&self.len.to_string());
}
TemplatePart::BytesPerSecond => {
result.push_str(&format!("{}/s", bytes_to_human(bytes_per_second as u64)));
}
TemplatePart::Eta => {
if self.pos == 0 {
result.push_str("Unknown");
} else {
let eta = (self.len - self.pos) as f64 / bytes_per_second;
result.push_str(&duration_to_human(std::time::Duration::from_secs(
eta as u64,
)));
}
}
TemplatePart::Bar(size) => {
let filled = (self.pos as f64 / self.len as f64 * *size as f64) as usize;
if *size >= filled {
let empty = *size - filled;
result.push('[');
for _ in 0..filled {
result.push('=');
}
for _ in 0..empty {
result.push(' ');
}
result.push(']');
} else {
let overflowed = filled - *size;
result.push('[');
for _ in 0..*size {
result.push('=');
}
for _ in 0..overflowed {
result.push('!');
}
}
}
TemplatePart::StateEmoji => {
if self.pos == self.len {
result.push('✅');
} else if self.pos == 0 {
result.push('🆕');
} else if self.pos > self.len {
result.push('💥');
} else {
result.push('⏳');
}
}
}
}
result
}
}
pub struct Bar {
id: usize,
manager: Weak<ManagerInner>,
}
pub(crate) struct ManagerInner {
states: Mutex<BTreeMap<usize, Arc<Mutex<BarState>>>>,
ansi: Mutex<Option<bool>>,
interval: std::time::Duration,
pub(crate) out: Arc<Mutex<Box<dyn Out>>>,
ticker: Mutex<Option<Ticker>>,
force_when_finished: AtomicBool,
next_id: AtomicUsize,
last_draw: Mutex<std::time::Instant>,
last_lines: AtomicUsize,
need_redraw: AtomicBool,
}
impl ManagerInner {
pub(crate) fn is_ticker_enabled(&self) -> bool {
self.ticker.lock().unwrap().is_some()
}
pub(crate) fn clear_existing(&self, out: &mut Box<dyn Out>) {
for _ in 0..self.last_lines.load(std::sync::atomic::Ordering::Acquire) {
let _ = out.write_all(format!("{}{}", UP_ANSI, CLEAR_ANSI).as_bytes());
}
}
pub(crate) fn is_terminal(&self, out: &mut Box<dyn Out>) -> bool {
let ansi = self.ansi.lock().unwrap();
match *ansi {
None => out.is_terminal(),
Some(force) => force,
}
}
pub(crate) fn draw_inner(
&self,
states: &BTreeMap<usize, Arc<Mutex<BarState>>>,
out: &mut Box<dyn Out>,
is_terminal: bool,
) {
let mut newlines = 0;
for state in states.values() {
let mut state = state.lock().unwrap();
if !state.visible {
continue;
}
if !is_terminal && !state.need_redraw {
continue;
}
let outstr = format!("{}\n", state.render());
if is_terminal {
let splits = outstr.split('\n');
let term_col = get_width(out.as_ref()) as usize;
for i in splits {
let width = string_width(i);
newlines += width / term_col;
if width % term_col != 0 {
newlines += 1;
}
}
}
let _ = out.write_all(outstr.as_bytes());
state.need_redraw = false;
}
if is_terminal {
self.last_lines
.store(newlines, std::sync::atomic::Ordering::Release);
}
}
pub(crate) fn mark_redraw(&self) {
self.need_redraw
.store(true, std::sync::atomic::Ordering::Release);
}
pub(crate) fn draw(&self, force: bool) {
if !force && self.is_ticker_enabled() {
return;
}
let now = std::time::Instant::now();
let mut last_draw = self.last_draw.lock().unwrap();
if !force && now - *last_draw < self.interval {
return;
}
if !self
.need_redraw
.swap(false, std::sync::atomic::Ordering::AcqRel)
{
return;
}
let mut out = self.out.lock().unwrap();
let states = self.states.lock().unwrap();
let is_terminal = self.is_terminal(&mut out);
if is_terminal && states.len() > 0 {
self.clear_existing(&mut out);
}
self.draw_inner(&states, &mut out, is_terminal);
*last_draw = now;
}
pub(crate) fn suspend<F: FnOnce(&mut Box<dyn Out>) -> R, R>(&self, f: F) -> R {
let mut out = self.out.lock().unwrap();
let is_terminal = self.is_terminal(&mut out);
if is_terminal {
self.clear_existing(&mut out);
}
let result = f(&mut out);
if is_terminal {
let states = self.states.lock().unwrap();
self.draw_inner(&states, &mut out, is_terminal);
}
result
}
}
#[cfg(all(unix, feature = "console_width"))]
pub trait Out: std::io::Write + std::io::IsTerminal + std::os::fd::AsRawFd + Send + Sync {}
#[cfg(all(unix, feature = "console_width"))]
impl<T: std::io::Write + std::io::IsTerminal + std::os::fd::AsRawFd + Send + Sync> Out for T {}
#[cfg(all(windows, feature = "console_width"))]
pub trait Out:
std::io::Write + std::io::IsTerminal + std::os::windows::io::AsRawHandle + Send + Sync
{
}
#[cfg(all(windows, feature = "console_width"))]
impl<T: std::io::Write + std::io::IsTerminal + std::os::windows::io::AsRawHandle + Send + Sync> Out
for T
{
}
#[cfg(not(any(
all(windows, feature = "console_width"),
all(unix, feature = "console_width")
)))]
pub trait Out: std::io::Write + std::io::IsTerminal + Send + Sync {}
#[cfg(not(any(
all(windows, feature = "console_width"),
all(unix, feature = "console_width")
)))]
impl<T: std::io::Write + std::io::IsTerminal + Send + Sync> Out for T {}
pub struct Manager {
inner: Arc<ManagerInner>,
}
impl Manager {
pub fn new(interval: std::time::Duration) -> Self {
Manager {
inner: Arc::new(ManagerInner {
states: Mutex::new(BTreeMap::new()),
next_id: AtomicUsize::new(0),
interval,
out: Arc::new(Mutex::new(Box::new(std::io::stdout()))),
last_draw: Mutex::new(std::time::Instant::now() - interval),
last_lines: AtomicUsize::new(0),
ansi: Mutex::new(None),
need_redraw: AtomicBool::new(false),
ticker: Mutex::new(None),
force_when_finished: AtomicBool::new(true),
}),
}
}
fn mark_redraw(&self) {
self.inner.mark_redraw();
}
pub fn with_stdout(self) -> Self {
*self.inner.out.lock().unwrap() = Box::new(std::io::stdout());
self.mark_redraw();
self
}
pub fn with_stderr(self) -> Self {
*self.inner.out.lock().unwrap() = Box::new(std::io::stderr());
self.mark_redraw();
self
}
pub fn with_file(self, file: std::fs::File) -> Self {
*self.inner.out.lock().unwrap() = Box::new(file);
self.mark_redraw();
self
}
pub fn auto_ansi(self) -> Self {
*self.inner.ansi.lock().unwrap() = None;
self.mark_redraw();
self
}
pub fn force_ansi(self, force: bool) -> Self {
*self.inner.ansi.lock().unwrap() = Some(force);
self.mark_redraw();
self
}
pub fn set_ticker(&self, set_ticker: bool) {
let mut ticker = self.inner.ticker.lock().unwrap();
if set_ticker && ticker.is_none() {
*ticker = Some(Ticker::new(self.inner.clone()));
} else if !set_ticker && ticker.is_some() {
*ticker = None;
}
}
pub fn force_draw_when_finished(&self, force: bool) {
self.inner
.force_when_finished
.store(force, std::sync::atomic::Ordering::Release);
}
pub fn create_bar(&self, len: u64, message: &str, template: &str, visible: bool) -> Bar {
let id = self
.inner
.next_id
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let bar_state = Arc::new(Mutex::new(BarState {
len,
pos: 0,
message: message.to_string(),
template: Template::new(template),
created_at: std::time::Instant::now(),
visible,
need_redraw: true,
}));
self.inner
.states
.lock()
.unwrap()
.insert(id, bar_state.clone());
if visible {
self.mark_redraw();
self.draw(true);
}
Bar {
manager: Arc::downgrade(&self.inner),
id,
}
}
pub fn draw(&self, force: bool) {
self.inner.draw(force);
}
pub fn suspend<F: FnOnce(&mut Box<dyn Out>) -> R, R>(&self, f: F) -> R {
self.inner.suspend(f)
}
pub fn create_writer(&self) -> writer::KyuriWriter {
writer::KyuriWriter::new(self.inner.clone())
}
}
impl Drop for ManagerInner {
fn drop(&mut self) {
self.draw(true);
}
}
impl Bar {
fn get_manager_and_state(&self) -> Option<(Arc<ManagerInner>, Arc<Mutex<BarState>>)> {
let manager = self.manager.upgrade()?;
let state = manager.states.lock().unwrap().get(&self.id)?.clone();
Some((manager, state))
}
fn check_if_force_draw(&self, manager: Arc<ManagerInner>, pos: u64, len: u64) {
if pos == len && manager
.force_when_finished
.load(std::sync::atomic::Ordering::Acquire)
{
manager.draw(true);
} else {
manager.draw(false);
}
}
pub fn inc(&self, n: u64) {
if let Some((manager, state)) = self.get_manager_and_state() {
let mut state = state.lock().unwrap();
state.pos += n;
state.need_redraw = true;
let pos = state.pos;
let len = state.len;
std::mem::drop(state);
manager.mark_redraw();
self.check_if_force_draw(manager, pos, len);
}
}
pub fn set_pos(&self, pos: u64) {
if let Some((manager, state)) = self.get_manager_and_state() {
let mut state = state.lock().unwrap();
state.pos = pos;
state.need_redraw = true;
let pos = state.pos;
let len = state.len;
std::mem::drop(state);
manager.mark_redraw();
self.check_if_force_draw(manager, pos, len);
}
}
pub fn set_len(&self, len: u64) {
if let Some((manager, state)) = self.get_manager_and_state() {
let mut state = state.lock().unwrap();
state.len = len;
state.need_redraw = true;
let pos = state.pos;
let len = state.len;
std::mem::drop(state);
manager.mark_redraw();
self.check_if_force_draw(manager, pos, len);
}
}
pub fn reset_created_at(&self) {
if let Some((manager, state)) = self.get_manager_and_state() {
let mut state = state.lock().unwrap();
state.created_at = std::time::Instant::now();
state.need_redraw = true;
let pos = state.pos;
let len = state.len;
std::mem::drop(state);
manager.mark_redraw();
self.check_if_force_draw(manager, pos, len);
}
}
pub fn get_pos(&self) -> u64 {
self.get_manager_and_state()
.map_or(0, |(_, state)| state.lock().unwrap().pos)
}
pub fn get_len(&self) -> u64 {
self.get_manager_and_state()
.map_or(0, |(_, state)| state.lock().unwrap().len)
}
pub fn finish(&self) {
if let Some((manager, state)) = self.get_manager_and_state() {
let state = state.lock().unwrap();
let pos = state.pos;
let len = state.len;
if pos != len {
self.set_pos(len);
}
std::mem::drop(state);
manager.draw(true);
}
}
pub fn finish_and_drop(self) {
self.finish();
}
pub fn set_visible(&self, visible: bool) {
if let Some((manager, state)) = self.get_manager_and_state() {
let mut state = state.lock().unwrap();
if state.visible != visible {
state.visible = visible;
state.need_redraw = true;
std::mem::drop(state);
manager.mark_redraw();
manager.draw(true);
}
}
}
pub fn is_visible(&self) -> bool {
self.get_manager_and_state()
.map_or(false, |(_, state)| state.lock().unwrap().visible)
}
pub fn set_message(&self, message: &str) {
if let Some((manager, state)) = self.get_manager_and_state() {
let mut state = state.lock().unwrap();
state.message = message.to_string();
state.need_redraw = true;
std::mem::drop(state);
manager.mark_redraw();
manager.draw(false);
}
}
pub fn set_template(&self, template: &str) {
if let Some((manager, state)) = self.get_manager_and_state() {
let mut state = state.lock().unwrap();
state.template = Template::new(template);
state.need_redraw = true;
std::mem::drop(state);
manager.mark_redraw();
manager.draw(false);
}
}
pub fn alive(&self) -> bool {
self.get_manager_and_state().is_some()
}
}
impl Drop for Bar {
fn drop(&mut self) {
if let Some((manager, _)) = self.get_manager_and_state() {
manager.states.lock().unwrap().remove(&self.id);
manager.mark_redraw();
manager.draw(true);
}
}
}
#[cfg(test)]
mod tests {
use std::io::{Read, Seek};
use super::*;
#[test]
fn basic_test() {
let manager = Manager::new(std::time::Duration::from_secs(1));
let bar_1 = manager.create_bar(
100,
"Downloading",
"{msg}\n[{elapsed}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
true,
);
let bar_2 = manager.create_bar(
100,
"Uploading",
"{msg}\n[{elapsed}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
true,
);
bar_1.set_pos(50);
bar_2.set_pos(25);
std::mem::drop(bar_1);
std::mem::drop(bar_2);
}
#[test]
fn dont_crash_when_zero() {
let manager = Manager::new(std::time::Duration::from_secs(1));
let bar = manager.create_bar(
0,
"Downloading",
"{msg}\n[{elapsed}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
true,
);
bar.set_pos(0);
manager.draw(true);
}
#[test]
fn inc() {
let manager = Manager::new(std::time::Duration::from_secs(1));
let bar = manager.create_bar(
100,
"Downloading",
"{msg}\n[{elapsed}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
true,
);
bar.inc(10);
bar.inc(10);
bar.inc(10);
bar.inc(10);
bar.inc(10);
assert_eq!(bar.get_pos(), 50);
std::mem::drop(bar);
}
#[test]
fn visible() {
let manager = Manager::new(std::time::Duration::from_secs(1));
let bar = manager.create_bar(
100,
"Downloading",
"{msg}\n[{elapsed}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
true,
);
assert!(bar.is_visible());
bar.set_visible(false);
assert!(!bar.is_visible());
std::mem::drop(bar);
}
#[test]
fn ticker() {
let manager = Manager::new(std::time::Duration::from_secs(1));
manager.set_ticker(true);
let bar = manager.create_bar(
100,
"Downloading",
"{msg}\n[{elapsed}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
true,
);
std::thread::sleep(std::time::Duration::from_secs(2));
std::mem::drop(bar);
}
#[test]
fn alive() {
let manager = Manager::new(std::time::Duration::from_secs(1));
let bar = manager.create_bar(
100,
"Downloading",
"{msg}\n[{elapsed}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
true,
);
assert!(bar.alive());
std::mem::drop(manager);
assert!(!bar.alive());
}
#[cfg(target_os = "linux")]
#[test]
fn test_pb_to_file() {
const TEMPLATE_SIMPLE: &str = "{msg}\n{bytes}/{total_bytes}";
let memfd_name = std::ffi::CString::new("test_pb_to_file").unwrap();
let memfd_fd =
nix::sys::memfd::memfd_create(&memfd_name, nix::sys::memfd::MemFdCreateFlag::empty())
.unwrap();
let memfd_writer: std::fs::File = memfd_fd.into();
let mut memfd_writer_clone = memfd_writer.try_clone().unwrap();
let progressbar_manager =
Manager::new(std::time::Duration::from_secs(1)).with_file(memfd_writer);
let pb1 = progressbar_manager.create_bar(
10,
"Downloading http://d1.example.com/",
TEMPLATE_SIMPLE,
true,
);
let pb2 = progressbar_manager.create_bar(
10,
"Downloading http://d2.example.com/",
TEMPLATE_SIMPLE,
true,
);
pb1.set_pos(2);
pb2.set_pos(3);
progressbar_manager.draw(true);
pb1.set_pos(5);
pb2.set_pos(7);
std::mem::drop(progressbar_manager);
memfd_writer_clone
.seek(std::io::SeekFrom::Start(0))
.unwrap();
let mut output = String::new();
memfd_writer_clone.read_to_string(&mut output).unwrap();
assert_eq!(
output,
r#"Downloading http://d1.example.com/
0 B/10 B
Downloading http://d2.example.com/
0 B/10 B
Downloading http://d1.example.com/
2 B/10 B
Downloading http://d2.example.com/
3 B/10 B
Downloading http://d1.example.com/
5 B/10 B
Downloading http://d2.example.com/
7 B/10 B
"#
);
}
}