pistonite_cu/cli/
flags.rs

1use std::ffi::OsString;
2use std::sync::Arc;
3use std::time::Instant;
4
5use clap::{Command, CommandFactory, FromArgMatches, Parser};
6
7use crate::cli::LogConfig;
8use crate::lv;
9
10/// Common flags for `cu::cli`
11#[derive(Default, Debug, Clone, PartialEq, Parser)]
12pub struct Flags {
13    /// Verbose. More -v makes it more verbose (opposite of --quiet)
14    #[clap(short = 'v', long, action(clap::ArgAction::Count))]
15    pub verbose: u8,
16    /// Quiet. More -q makes it more quiet (opposite of --verbose)
17    #[clap(short = 'q', long, action(clap::ArgAction::Count))]
18    pub quiet: u8,
19    /// Set the color mode for this program. May affect subprocesses spawned.
20    #[clap(long)]
21    pub color: Option<lv::Color>,
22    /// Automatically answer 'yes' to all yes/no prompts
23    #[cfg(feature = "prompt")]
24    #[clap(short = 'y', long)]
25    pub yes: bool,
26    /// Make all prompts fail with an error. (Cancels with one --interactive)
27    #[cfg(feature = "prompt")]
28    #[clap(long, action(clap::ArgAction::Count))]
29    pub non_interactive: u8,
30    /// Allow interactivity. Cancels with one --non-interactive
31    #[cfg(feature = "prompt")]
32    #[clap(long, action(clap::ArgAction::Count))]
33    pub interactive: u8,
34}
35
36impl AsRef<Flags> for Flags {
37    fn as_ref(&self) -> &Flags {
38        self
39    }
40}
41
42impl Flags {
43    /// Apply the CLI Flags
44    ///
45    /// # Safety
46    /// This is unsafe because it modifies environment variables.
47    /// The [`cu::cli`](macro@crate::cli) macro generates safe call to this
48    /// when the program only has the main thread.
49    pub unsafe fn apply(&self, log_config: Arc<dyn LogConfig + Send + Sync>) {
50        let level = self.print_level();
51        if level == lv::Print::VerboseVerbose {
52            if std::env::var("RUST_BACKTRACE")
53                .unwrap_or_default()
54                .is_empty()
55            {
56                unsafe { std::env::set_var("RUST_BACKTRACE", "1") }
57            }
58        }
59
60        let prompt = {
61            #[cfg(feature = "prompt")]
62            match self.non_interactive.min(i8::MAX as u8) as i8
63                - self.interactive.min(i8::MAX as u8) as i8
64            {
65                ..=0 => {
66                    if self.yes {
67                        Some(lv::Prompt::YesOrInteractive)
68                    } else {
69                        Some(lv::Prompt::Interactive)
70                    }
71                }
72                _ => {
73                    if self.yes {
74                        Some(lv::Prompt::YesOrBlock)
75                    } else {
76                        Some(lv::Prompt::Block)
77                    }
78                }
79            }
80            #[cfg(not(feature = "prompt"))]
81            {
82                None
83            }
84        };
85        super::print_init::init_options(self.color.unwrap_or_default(), level, prompt, log_config);
86    }
87
88    pub fn print_level(&self) -> lv::Print {
89        let level = (self.verbose as i8 - self.quiet as i8).clamp(-2, 2);
90        level.into()
91    }
92
93    /// Merge `other` into self. Options in other will be applied on top of self (equivalent
94    /// to specifying `self` then specify `other`
95    pub fn merge(&mut self, other: &Self) {
96        self.verbose += other.verbose;
97        self.quiet += other.quiet;
98        if let Some(color) = other.color {
99            self.color = Some(color);
100        }
101        #[cfg(feature = "prompt")]
102        {
103            if other.yes {
104                self.yes = true;
105            }
106            self.non_interactive += other.non_interactive;
107            self.interactive += other.interactive;
108        }
109    }
110}
111
112/// Entry point to CLI
113///
114/// # Safety
115/// A safe wrapper is generated by the [`cu::cli`](macro@crate::cli) macro.
116/// See [module level documentation](self) for more.
117#[inline(always)]
118#[doc(hidden)]
119pub unsafe fn __run<
120    TArg: clap::Parser,
121    TLogConfig: LogConfig + Send + Sync + 'static,
122    FPreproc: FnOnce(&mut TArg),
123    FLogConfig: FnOnce(&Flags) -> TLogConfig,
124    FExecute: FnOnce(TArg) -> crate::Result<()>,
125    FFlag: FnOnce(&TArg) -> &Flags,
126>(
127    fn_preproc: FPreproc,
128    fn_log_config: FLogConfig,
129    fn_execute: FExecute,
130    fn_flag: FFlag,
131) -> std::process::ExitCode {
132    let start = std::time::Instant::now();
133    let args = unsafe {
134        parse_args_or_help::<TArg, TLogConfig, FPreproc, FLogConfig, FFlag>(
135            fn_preproc,
136            fn_log_config,
137            fn_flag,
138        )
139    };
140    let result = fn_execute(args);
141    handle_result(start, result)
142}
143
144/// Entry point to CLI
145///
146/// # Safety
147/// A safe wrapper is generated by the [`cu::cli`](macro@crate::cli) macro.
148/// See [module level documentation](self) for more.
149#[inline(always)]
150#[cfg(feature = "coroutine")]
151#[doc(hidden)]
152pub unsafe fn __co_run<
153    TArg: clap::Parser + Send + 'static,
154    TLogConfig: LogConfig + Send + Sync + 'static,
155    FPreproc: FnOnce(&mut TArg),
156    FLogConfig: FnOnce(&Flags) -> TLogConfig,
157    FExecute: FnOnce(TArg) -> TResult + Send + 'static,
158    TResult: Future<Output = crate::Result<()>> + Send + 'static,
159    FFlag: FnOnce(&TArg) -> &Flags,
160>(
161    fn_preproc: FPreproc,
162    fn_log_config: FLogConfig,
163    fn_execute: FExecute,
164    fn_flag: FFlag,
165) -> std::process::ExitCode {
166    let start = std::time::Instant::now();
167    let args = unsafe {
168        parse_args_or_help::<TArg, TLogConfig, FPreproc, FLogConfig, FFlag>(
169            fn_preproc,
170            fn_log_config,
171            fn_flag,
172        )
173    };
174    #[cfg(not(feature = "coroutine-heavy"))]
175    let result = crate::co::block(async move { fn_execute(args).await });
176    #[cfg(feature = "coroutine-heavy")]
177    let result = crate::co::run(async move { fn_execute(args).await });
178
179    handle_result(start, result)
180}
181
182unsafe fn parse_args_or_help<
183    TArg: Parser,
184    TLogConfig: LogConfig + Send + Sync + 'static,
185    FPreproc: FnOnce(&mut TArg),
186    FLogConfig: FnOnce(&Flags) -> TLogConfig,
187    FFlag: FnOnce(&TArg) -> &Flags,
188>(
189    fn_preproc: FPreproc,
190    fn_log_config: FLogConfig,
191    fn_flag: FFlag,
192) -> TArg {
193    let mut parsed = parse_args::<TArg>();
194    fn_preproc(&mut parsed);
195    let flags = fn_flag(&parsed);
196    let log_config: Arc<dyn LogConfig + Send + Sync> = Arc::new(fn_log_config(flags));
197    unsafe { flags.apply(log_config) };
198    parsed
199}
200
201/// Wrapper for clap parse to respect the color flag when printing help or error
202fn parse_args<T: Parser>() -> T {
203    // parse the color arg first, so that we can respect it when printing help
204    let color = lv::Color::from_os_args();
205    let use_color = color.is_colored_for_stdout();
206
207    // this will exit on error
208    let mut matches = get_colored_command::<T>(use_color).get_matches();
209
210    match <T as FromArgMatches>::from_arg_matches_mut(&mut matches) {
211        Ok(x) => x,
212        Err(e) => {
213            let mut command = get_colored_command::<T>(use_color);
214            let error = e.format(&mut command);
215            error.exit()
216        }
217    }
218}
219
220/// Try to parse arguments from an iterator and print the error/help
221/// on failure.
222///
223/// Whether the output has color depends on the main CLI `--color` option.
224/// This is useful for implementing custom command parser within
225/// an application.
226pub fn try_parse<T: Parser, I: IntoIterator>(iter: I) -> Option<T>
227where
228    I::Item: Into<OsString> + Clone,
229{
230    let use_color = crate::lv::color_enabled();
231    let result = get_colored_command::<T>(use_color)
232        .try_get_matches_from(iter)
233        .and_then(|mut matches| <T as FromArgMatches>::from_arg_matches_mut(&mut matches));
234    match result {
235        Ok(x) => Some(x),
236        Err(e) => {
237            let mut command = get_colored_command::<T>(use_color);
238            let error = e.format(&mut command);
239            if let Err(e) = error.print() {
240                crate::warn!("arg parse error failed to print: {e:?}");
241            }
242            None
243        }
244    }
245}
246
247/// Print the help text from a command parser.
248///
249/// Whether the output has color depends on the main CLI `--color` option.
250/// This is useful for implementing custom command parser within
251/// an application.
252#[inline(always)]
253pub fn print_help<T: Parser>(long: bool) {
254    let command = get_colored_command::<T>(crate::lv::color_enabled());
255    print_help_impl(command, long)
256}
257fn print_help_impl(mut command: Command, long: bool) {
258    let result = if long {
259        command.print_long_help()
260    } else {
261        command.print_help()
262    };
263    if let Err(e) = result {
264        crate::warn!("help failed to print: {e:?}");
265    }
266}
267
268#[inline(always)]
269fn get_colored_command<T: Parser>(color: bool) -> Command {
270    get_colored_command_impl(<T as CommandFactory>::command(), color)
271}
272fn get_colored_command_impl(command: Command, color: bool) -> Command {
273    use clap::builder::styling::{AnsiColor, Styles};
274    if color {
275        // Modified version of Cargo's color style
276        // [source](https://github.com/crate-ci/clap-cargo/blob/master/src/style.rs)
277        command.styles(
278            Styles::styled()
279                .header(AnsiColor::BrightYellow.on_default())
280                .usage(AnsiColor::BrightRed.on_default())
281                .literal(AnsiColor::BrightCyan.on_default())
282                .placeholder(AnsiColor::Cyan.on_default())
283                .error(AnsiColor::BrightRed.on_default())
284                .valid(AnsiColor::BrightCyan.on_default())
285                .invalid(AnsiColor::BrightYellow.on_default())
286                .context(AnsiColor::BrightYellow.on_default()),
287        )
288    } else {
289        command.styles(Styles::plain())
290    }
291}
292
293fn handle_result(start: Instant, result: crate::Result<()>) -> std::process::ExitCode {
294    let elapsed = start.elapsed().as_secs_f32();
295    if let Err(e) = result {
296        crate::error!("fatal: {e:?}");
297        // we display the hint for user to use -vv
298        // if:
299        // - the user is already tried to get more debug info with -v
300        //   (because otherwise it will be too noisy and it might be a user-error,
301        //   not a bug
302        // - the trace hint is not explicitly disabled
303        // - the trace hint is not already displayed
304        if lv::D.enabled() && crate::lv::is_trace_hint_enabled() {
305            if std::env::var("RUST_BACKTRACE")
306                .unwrap_or_default()
307                .is_empty()
308            {
309                crate::hint!("use -vv or set RUST_BACKTRACE=1 to display the error backtrace.");
310            }
311        }
312        if crate::lv::is_print_time_enabled() {
313            // use debug so the error trace is the last line,
314            // so user is directed to see what is the most important
315            crate::debug!("finished in {elapsed:.2}s");
316        }
317        reset_color();
318        std::process::ExitCode::FAILURE
319    } else {
320        if crate::lv::is_print_time_enabled() {
321            crate::info!("finished in {elapsed:.2}s");
322        }
323        reset_color();
324        std::process::ExitCode::SUCCESS
325    }
326}
327
328/// Reset terminal color
329///
330/// **You do not need to call this if using the `#[cu::cli]` macro**.
331pub fn reset_color() {
332    use std::io::IsTerminal as _;
333    use std::io::Write as _;
334    let mut stdout = std::io::stdout();
335    if stdout.is_terminal() {
336        let _ = write!(stdout, "\x1b[0m");
337        let _ = stdout.flush();
338        return;
339    }
340    let mut stderr = std::io::stderr();
341    if stderr.is_terminal() {
342        let _ = write!(stderr, "\x1b[0m");
343        let _ = stderr.flush();
344    }
345}
pistonite_cu/cli/flags.rs

pistonite_cu/cli/
flags.rs
