Skip to content

Commit

Permalink
Implement customizable tab expansion (console-rs#150)
Browse files Browse the repository at this point in the history
  • Loading branch information
chris-laplante committed Jul 6, 2022
1 parent 72c25c1 commit c9596f3
Show file tree
Hide file tree
Showing 3 changed files with 148 additions and 71 deletions.
69 changes: 31 additions & 38 deletions src/progress_bar.rs
Expand Up @@ -5,12 +5,11 @@ use std::sync::{Arc, Condvar, Mutex, MutexGuard, Weak};
use std::time::{Duration, Instant};
use std::{fmt, io, thread};

use console::strip_ansi_codes;
#[cfg(test)]
use once_cell::sync::Lazy;

use crate::draw_target::ProgressDrawTarget;
use crate::state::{AtomicPosition, BarState, ProgressFinish, Reset};
use crate::state::{AtomicPosition, BarState, ProgressFinish, Reset, TabExpandedString};
use crate::style::ProgressStyle;
use crate::{ProgressBarIter, ProgressIterator, ProgressState};

Expand Down Expand Up @@ -70,15 +69,25 @@ impl ProgressBar {
self
}

/// A convenience builder-like function for a progress bar with a given tab width
pub fn with_tab_width(self, tab_width: usize) -> ProgressBar {
self.state().set_tab_width(tab_width);
self
}

/// A convenience builder-like function for a progress bar with a given prefix
pub fn with_prefix(self, prefix: impl Into<Cow<'static, str>>) -> ProgressBar {
self.state().state.prefix = prefix.into();
let mut state = self.state();
state.state.prefix = TabExpandedString::new(prefix.into(), state.tab_width);
drop(state);
self
}

/// A convenience builder-like function for a progress bar with a given message
pub fn with_message(self, message: impl Into<Cow<'static, str>>) -> ProgressBar {
self.state().state.message = message.into();
let mut state = self.state();
state.state.message = TabExpandedString::new(message.into(), state.tab_width);
drop(state);
self
}

Expand Down Expand Up @@ -123,7 +132,14 @@ impl ProgressBar {
///
/// This does not redraw the bar. Call [`ProgressBar::tick()`] to force it.
pub fn set_style(&self, style: ProgressStyle) {
self.state().style = style;
self.state().set_style(style);
}

/// Sets the tab width (default: 8). All tabs will be expanded to this many spaces.
pub fn set_tab_width(&mut self, tab_width: usize) {
let mut state = self.state();
state.set_tab_width(tab_width);
state.draw(true, Instant::now()).unwrap();
}

/// Spawns a background thread to tick the progress bar
Expand Down Expand Up @@ -244,15 +260,19 @@ impl ProgressBar {
/// For the prefix to be visible, the `{prefix}` placeholder must be present in the template
/// (see [`ProgressStyle`]).
pub fn set_prefix(&self, prefix: impl Into<Cow<'static, str>>) {
self.state().set_prefix(Instant::now(), prefix.into());
let mut state = self.state();
state.state.prefix = TabExpandedString::new(prefix.into(), state.tab_width);
state.update_estimate_and_draw(Instant::now());
}

/// Sets the current message of the progress bar
///
/// For the message to be visible, the `{msg}` placeholder must be present in the template (see
/// [`ProgressStyle`]).
pub fn set_message(&self, msg: impl Into<Cow<'static, str>>) {
self.state().set_message(Instant::now(), msg.into())
let mut state = self.state();
state.state.message = TabExpandedString::new(msg.into(), state.tab_width);
state.update_estimate_and_draw(Instant::now());
}

/// Creates a new weak reference to this `ProgressBar`
Expand Down Expand Up @@ -516,23 +536,13 @@ impl ProgressBar {
}

/// Current message
pub fn message(&self) -> Cow<'static, str> {
self.state().state.message.clone()
}

/// Current message (with ANSI escape codes stripped)
pub fn message_unstyled(&self) -> String {
strip_ansi_codes(&self.message()).to_string()
pub fn message(&self) -> String {
self.state().state.message.expanded().to_string()
}

/// Current prefix
pub fn prefix(&self) -> Cow<'static, str> {
self.state().state.prefix.clone()
}

/// Current prefix (with ANSI escape codes stripped)
pub fn prefix_unstyled(&self) -> String {
strip_ansi_codes(&self.prefix()).to_string()
pub fn prefix(&self) -> String {
self.state().state.prefix.expanded().to_string()
}

#[inline]
Expand Down Expand Up @@ -662,7 +672,6 @@ pub(crate) static TICKER_TEST: Lazy<Mutex<()>> = Lazy::new(Mutex::default);
#[cfg(test)]
mod tests {
use super::*;
use console::Style;

#[allow(clippy::float_cmp)]
#[test]
Expand Down Expand Up @@ -762,20 +771,4 @@ mod tests {
drop(pb2);
assert!(!TICKER_RUNNING.load(Ordering::SeqCst));
}

#[test]
fn access_message_and_prefix() {
let pb = ProgressBar::new(80);
pb.set_message(Style::new().red().bold().apply_to("text").to_string());
pb.set_prefix(
Style::new()
.on_blue()
.italic()
.apply_to("prefix!")
.to_string(),
);

assert_eq!(pb.message_unstyled(), "text");
assert_eq!(pb.prefix_unstyled(), "prefix!");
}
}
83 changes: 70 additions & 13 deletions src/state.rs
Expand Up @@ -12,6 +12,7 @@ pub(crate) struct BarState {
pub(crate) on_finish: ProgressFinish,
pub(crate) style: ProgressStyle,
pub(crate) state: ProgressState,
pub(crate) tab_width: usize,
}

impl BarState {
Expand All @@ -25,6 +26,7 @@ impl BarState {
on_finish: ProgressFinish::default(),
style: ProgressStyle::default_bar(),
state: ProgressState::new(len, pos),
tab_width: DEFAULT_TAB_WIDTH,
}
}

Expand All @@ -42,7 +44,7 @@ impl BarState {
if let Some(len) = self.state.len {
self.state.pos.set(len);
}
self.state.message = msg;
self.state.message = TabExpandedString::new(msg, self.tab_width);
}
ProgressFinish::AndClear => {
if let Some(len) = self.state.len {
Expand All @@ -51,7 +53,9 @@ impl BarState {
self.state.status = Status::DoneHidden;
}
ProgressFinish::Abandon => {}
ProgressFinish::AbandonWithMessage(msg) => self.state.message = msg,
ProgressFinish::AbandonWithMessage(msg) => {
self.state.message = TabExpandedString::new(msg, self.tab_width)
}
}

// There's no need to update the estimate here; once the `status` is no longer
Expand Down Expand Up @@ -92,22 +96,24 @@ impl BarState {
self.update_estimate_and_draw(now);
}

pub(crate) fn set_message(&mut self, now: Instant, msg: Cow<'static, str>) {
self.state.message = msg;
self.update_estimate_and_draw(now);
pub(crate) fn set_tab_width(&mut self, tab_width: usize) {
self.tab_width = tab_width;
self.state.message.set_tab_width(tab_width);
self.state.prefix.set_tab_width(tab_width);
self.style.set_tab_width(tab_width);
}

pub(crate) fn set_prefix(&mut self, now: Instant, prefix: Cow<'static, str>) {
self.state.prefix = prefix;
self.update_estimate_and_draw(now);
pub(crate) fn set_style(&mut self, style: ProgressStyle) {
self.style = style;
self.style.set_tab_width(self.tab_width);
}

pub(crate) fn tick(&mut self, now: Instant) {
self.state.tick = self.state.tick.saturating_add(1);
self.update_estimate_and_draw(now);
}

fn update_estimate_and_draw(&mut self, now: Instant) {
pub(crate) fn update_estimate_and_draw(&mut self, now: Instant) {
let pos = self.state.pos.pos.load(Ordering::Relaxed);
self.state.est.record(pos, now);
let _ = self.draw(false, now);
Expand Down Expand Up @@ -190,8 +196,8 @@ pub struct ProgressState {
pub(crate) started: Instant,
status: Status,
est: Estimator,
pub(crate) message: Cow<'static, str>,
pub(crate) prefix: Cow<'static, str>,
pub(crate) message: TabExpandedString,
pub(crate) prefix: TabExpandedString,
}

impl ProgressState {
Expand All @@ -203,8 +209,8 @@ impl ProgressState {
status: Status::InProgress,
started: Instant::now(),
est: Estimator::new(Instant::now()),
message: "".into(),
prefix: "".into(),
message: TabExpandedString::NoTabs("".into()),
prefix: TabExpandedString::NoTabs("".into()),
}
}

Expand Down Expand Up @@ -289,6 +295,55 @@ impl ProgressState {
}
}

#[derive(Debug, PartialEq, Eq, Clone)]
pub(crate) enum TabExpandedString {
NoTabs(Cow<'static, str>),
WithTabs {
original: Cow<'static, str>,
expanded: String,
tab_width: usize,
},
}

impl TabExpandedString {
pub(crate) fn new(s: Cow<'static, str>, tab_width: usize) -> Self {
let expanded = s.replace('\t', &" ".repeat(tab_width));
if s == expanded {
Self::NoTabs(s)
} else {
Self::WithTabs {
original: s,
expanded,
tab_width,
}
}
}

pub(crate) fn expanded(&self) -> &str {
match &self {
Self::NoTabs(s) => {
debug_assert!(!s.contains('\t'));
s
}
Self::WithTabs { expanded, .. } => expanded,
}
}

pub(crate) fn set_tab_width(&mut self, new_tab_width: usize) {
if let TabExpandedString::WithTabs {
original,
expanded,
tab_width,
} = self
{
if *tab_width != new_tab_width {
*tab_width = new_tab_width;
*expanded = original.replace('\t', &" ".repeat(new_tab_width));
}
}
}
}

/// Estimate the number of seconds per step
///
/// Ring buffer with constant capacity. Used by `ProgressBar`s to display `{eta}`,
Expand Down Expand Up @@ -486,6 +541,8 @@ pub(crate) enum Status {
DoneHidden,
}

pub(crate) const DEFAULT_TAB_WIDTH: usize = 8;

#[cfg(test)]
mod tests {
use super::*;
Expand Down

0 comments on commit c9596f3

Please sign in to comment.