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#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
28pub struct ConsoleConfig {
29 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}