Implement customizable tab expansion (console-rs#150)

chris-laplante · chris-laplante · commit 70ebcebf064a · 2022-07-05T13:28:48.000-04:00
diff --git a/src/progress_bar.rs b/src/progress_bar.rs
@@ -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};
 
@@ -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
     }
 
@@ -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
@@ -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`
@@ -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]
@@ -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]
@@ -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!");
-    }
 }
diff --git a/src/state.rs b/src/state.rs
@@ -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 {
@@ -25,6 +26,7 @@ impl BarState {
             on_finish: ProgressFinish::default(),
             style: ProgressStyle::default_bar(),
             state: ProgressState::new(len, pos),
+            tab_width: DEFAULT_TAB_WIDTH,
         }
     }
 
@@ -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 {
@@ -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
@@ -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.change_tab_width(tab_width);
+        self.state.prefix.change_tab_width(tab_width);
+        self.style.change_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.change_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);
@@ -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 {
@@ -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()),
         }
     }
 
@@ -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 change_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}`,
@@ -486,6 +541,8 @@ pub(crate) enum Status {
     DoneHidden,
 }
 
+pub(crate) const DEFAULT_TAB_WIDTH: usize = 8;
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/style.rs b/src/style.rs
