acick_util/
console.rs

1use std::env;
2use std::io::{self, BufRead as _, Write};
3
4use anyhow::Context as _;
5use console::Term;
6use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle};
7
8static PB_TICK_INTERVAL_MS: u64 = 50;
9static PB_TEMPL_COUNT: &str =
10    "{spinner:.green} {prefix} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} {per_sec} ETA {eta}";
11static PB_TEMPL_BYTES: &str =
12    "{spinner:.green} {prefix} [{elapsed_precise}] [{wide_bar:.cyan/blue}] \
13     {bytes:>9}/{total_bytes:>9} {bytes_per_sec:>11} ETA {eta:>3}";
14static PB_PROGRESS_CHARS: &str = "#>-";
15
16#[derive(Debug)]
17enum Inner {
18    Term(Term),
19    Buf {
20        input: io::BufReader<io::Cursor<String>>,
21        output: Vec<u8>,
22    },
23    Sink(io::Sink),
24}
25
26/// Config for console.
27#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
28pub struct ConsoleConfig {
29    /// If true, assumes yes and skips any confirmation.
30    pub assume_yes: bool,
31}
32
33#[derive(Debug)]
34pub struct Console {
35    inner: Inner,
36    conf: ConsoleConfig,
37}
38
39impl Console {
40    pub fn term(conf: ConsoleConfig) -> Self {
41        Self {
42            inner: Inner::Term(Term::stderr()),
43            conf,
44        }
45    }
46
47    pub fn buf(conf: ConsoleConfig) -> Self {
48        Self {
49            inner: Inner::Buf {
50                input: io::BufReader::new(io::Cursor::new(String::new())),
51                output: Vec::new(),
52            },
53            conf,
54        }
55    }
56
57    pub fn sink(conf: ConsoleConfig) -> Self {
58        Self {
59            inner: Inner::Sink(io::sink()),
60            conf,
61        }
62    }
63
64    #[cfg(test)]
65    fn write_input(&mut self, s: &str) {
66        if let Inner::Buf { ref mut input, .. } = self.inner {
67            input.get_mut().get_mut().push_str(s)
68        }
69    }
70
71    pub fn take_buf(self) -> Option<Vec<u8>> {
72        match self.inner {
73            Inner::Buf { output: buf, .. } => Some(buf),
74            _ => None,
75        }
76    }
77
78    pub fn take_output(self) -> crate::Result<String> {
79        self.take_buf()
80            .context("Could not take buf from console")
81            .and_then(|buf| Ok(String::from_utf8(buf)?))
82    }
83
84    #[inline]
85    fn as_mut_write(&mut self) -> &mut dyn Write {
86        match self.inner {
87            Inner::Term(ref mut w) => w,
88            Inner::Buf {
89                output: ref mut w, ..
90            } => w,
91            Inner::Sink(ref mut w) => w,
92        }
93    }
94
95    pub fn warn(&mut self, message: &str) -> io::Result<()> {
96        writeln!(self, "WARN: {}", message)
97    }
98
99    pub fn confirm(&mut self, message: &str, default: bool) -> io::Result<bool> {
100        if self.conf.assume_yes {
101            return Ok(true);
102        }
103
104        let prompt = format!("{} ({}) ", message, if default { "Y/n" } else { "y/N" });
105        let input = self.prompt_and_read(&prompt, false)?;
106        match input.to_lowercase().as_str() {
107            "y" | "yes" => Ok(true),
108            "n" | "no" => Ok(false),
109            _ => Ok(default),
110        }
111    }
112
113    pub fn get_env_or_prompt_and_read(
114        &mut self,
115        env_name: &str,
116        prompt: &str,
117        is_password: bool,
118    ) -> io::Result<String> {
119        if let Ok(val) = env::var(env_name) {
120            writeln!(
121                self,
122                "{}{:16} (read from env {})",
123                prompt,
124                if is_password { "********" } else { &val },
125                env_name
126            )?;
127            return Ok(val);
128        };
129        self.prompt_and_read(prompt, is_password)
130    }
131
132    fn read_user(&mut self, is_password: bool) -> io::Result<String> {
133        match self.inner {
134            Inner::Term(ref term) => {
135                if is_password {
136                    term.read_secure_line()
137                } else {
138                    term.read_line()
139                }
140            }
141            Inner::Buf { ref mut input, .. } => {
142                let mut buf = String::new();
143                input.read_line(&mut buf)?;
144                Ok(buf)
145            }
146            Inner::Sink(_) => Ok(String::from("")),
147        }
148    }
149
150    fn prompt(&mut self, prompt: &str) -> io::Result<()> {
151        write!(self, "{}", prompt)?;
152        self.flush()?;
153        Ok(())
154    }
155
156    fn prompt_and_read(&mut self, prompt: &str, is_password: bool) -> io::Result<String> {
157        self.prompt(prompt)?;
158        self.read_user(is_password)
159    }
160
161    pub fn build_pb_count(&self, len: u64) -> ProgressBar {
162        self.build_pb_with(len, PB_TEMPL_COUNT)
163    }
164
165    pub fn build_pb_bytes(&self, len: u64) -> ProgressBar {
166        self.build_pb_with(len, PB_TEMPL_BYTES)
167    }
168
169    fn build_pb_with(&self, len: u64, template: &str) -> ProgressBar {
170        let pb = ProgressBar::with_draw_target(len, self.to_pb_target());
171        let style = Self::pb_style_common().template(template);
172        pb.set_style(style);
173        pb.enable_steady_tick(PB_TICK_INTERVAL_MS);
174        pb
175    }
176
177    fn to_pb_target(&self) -> ProgressDrawTarget {
178        match &self.inner {
179            Inner::Term(term) => ProgressDrawTarget::to_term(term.clone(), None),
180            _ => ProgressDrawTarget::hidden(),
181        }
182    }
183
184    fn pb_style_common() -> ProgressStyle {
185        ProgressStyle::default_bar().progress_chars(PB_PROGRESS_CHARS)
186    }
187}
188
189impl Write for Console {
190    #[inline]
191    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
192        self.as_mut_write().write(buf)
193    }
194
195    #[inline]
196    fn flush(&mut self) -> io::Result<()> {
197        self.as_mut_write().flush()
198    }
199}
200
201macro_rules! def_color {
202    ($name:ident, $name_upper:ident, $style:expr) => {
203        ::lazy_static::lazy_static! {
204            static ref $name_upper: ::console::Style = {
205                use ::console::Style;
206                $style
207            };
208        }
209
210        pub fn $name<D>(val: D) -> ::console::StyledObject<D> {
211            $name_upper.apply_to(val)
212        }
213    };
214}
215
216pub use color_defs::*;
217
218#[cfg_attr(coverage, no_coverage)]
219mod color_defs {
220    def_color!(sty_none, STY_NONE, Style::new());
221    def_color!(sty_r, STY_R, Style::new().red());
222    def_color!(sty_g, STY_G, Style::new().green());
223    def_color!(sty_y, STY_Y, Style::new().yellow());
224    def_color!(sty_dim, STY_DIM, Style::new().dim());
225    def_color!(sty_r_under, STY_R_UNDER, Style::new().underlined().red());
226    def_color!(sty_g_under, STY_G_UNDER, Style::new().underlined().green());
227    def_color!(sty_y_under, STY_Y_UNDER, Style::new().underlined().yellow());
228    def_color!(sty_r_rev, STY_R_REV, Style::new().bold().reverse().red());
229    def_color!(sty_g_rev, STY_G_REV, Style::new().bold().reverse().green());
230    def_color!(sty_y_rev, STY_Y_REV, Style::new().bold().reverse().yellow());
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn test_warn() -> anyhow::Result<()> {
239        let conf = ConsoleConfig { assume_yes: true };
240        let mut cnsl = Console::buf(conf);
241        cnsl.warn("message")?;
242        let output_str = cnsl.take_output()?;
243        assert_eq!(output_str, "WARN: message\n");
244        Ok(())
245    }
246
247    #[test]
248    fn test_confirm() -> anyhow::Result<()> {
249        let tests = &[
250            (true, "", false, true),
251            (false, "y", false, true),
252            (false, "Y", false, true),
253            (false, "yes", false, true),
254            (false, "Yes", false, true),
255            (false, "n", true, false),
256            (false, "N", true, false),
257            (false, "no", true, false),
258            (false, "No", true, false),
259            (false, "hoge", true, true),
260            (false, "hoge", false, false),
261            (false, "", true, true),
262            (false, "", false, false),
263        ];
264        for (assume_yes, input, default, expected) in tests {
265            let conf = ConsoleConfig {
266                assume_yes: *assume_yes,
267            };
268            let mut cnsl = Console::buf(conf);
269            cnsl.write_input(input);
270            let actual = cnsl.confirm("message", *default).unwrap();
271            assert_eq!(actual, *expected);
272        }
273        Ok(())
274    }
275
276    #[test]
277    fn test_get_env_or_prompt_and_read() -> anyhow::Result<()> {
278        let cnsl_term = Console::term(ConsoleConfig::default());
279        let cnsl_buf_0 = Console::buf(ConsoleConfig::default());
280        let mut cnsl_buf_1 = Console::buf(ConsoleConfig::default());
281        cnsl_buf_1.write_input("test_input");
282        let cnsl_sink_0 = Console::sink(ConsoleConfig::default());
283        let cnsl_sink_1 = Console::sink(ConsoleConfig::default());
284        let env_name_exists = if cfg!(windows) { "APPDATA" } else { "HOME" };
285        let env_val: &str = &env::var(env_name_exists).unwrap();
286        let tests = &mut [
287            (cnsl_term, env_name_exists, env_val),
288            (cnsl_buf_0, env_name_exists, env_val),
289            (cnsl_buf_1, "ACICK_TEST_UNKNOWN_VAR", "test_input"),
290            (cnsl_sink_0, env_name_exists, env_val),
291            (cnsl_sink_1, "ACICK_TEST_UNKNOWN_VAR", ""),
292        ];
293
294        for (ref mut cnsl, env_name, expected) in tests {
295            let actual = cnsl.get_env_or_prompt_and_read(env_name, "prompt >", true)?;
296            assert_eq!(&actual, expected);
297        }
298        Ok(())
299    }
300}