pistonite-cu 0.8.3

Battery-included common utils to speed up development of rust tools
Documentation
use std::ffi::OsString;
use std::sync::Arc;
use std::time::Instant;

use clap::{Command, CommandFactory, FromArgMatches, Parser};

use crate::cli::LogConfig;
use crate::lv;

/// Common flags for `cu::cli`
#[derive(Default, Debug, Clone, PartialEq, Parser)]
pub struct Flags {
    /// Verbose. More -v makes it more verbose (opposite of --quiet)
    #[clap(short = 'v', long, action(clap::ArgAction::Count))]
    pub verbose: u8,
    /// Quiet. More -q makes it more quiet (opposite of --verbose)
    #[clap(short = 'q', long, action(clap::ArgAction::Count))]
    pub quiet: u8,
    /// Set the color mode for this program. May affect subprocesses spawned.
    #[clap(long)]
    pub color: Option<lv::Color>,
    /// Automatically answer 'yes' to all yes/no prompts
    #[cfg(feature = "prompt")]
    #[clap(short = 'y', long)]
    pub yes: bool,
    /// Make all prompts fail with an error. (Cancels with one --interactive)
    #[cfg(feature = "prompt")]
    #[clap(long, action(clap::ArgAction::Count))]
    pub non_interactive: u8,
    /// Allow interactivity. Cancels with one --non-interactive
    #[cfg(feature = "prompt")]
    #[clap(long, action(clap::ArgAction::Count))]
    pub interactive: u8,
}

impl AsRef<Flags> for Flags {
    fn as_ref(&self) -> &Flags {
        self
    }
}

impl Flags {
    /// Apply the CLI Flags
    ///
    /// # Safety
    /// This is unsafe because it modifies environment variables.
    /// The [`cu::cli`](macro@crate::cli) macro generates safe call to this
    /// when the program only has the main thread.
    pub unsafe fn apply(&self, log_config: Arc<dyn LogConfig + Send + Sync>) {
        let level = self.print_level();
        if level == lv::Print::VerboseVerbose {
            if std::env::var("RUST_BACKTRACE")
                .unwrap_or_default()
                .is_empty()
            {
                unsafe { std::env::set_var("RUST_BACKTRACE", "1") }
            }
        }

        let prompt = {
            #[cfg(feature = "prompt")]
            match self.non_interactive.min(i8::MAX as u8) as i8
                - self.interactive.min(i8::MAX as u8) as i8
            {
                ..=0 => {
                    if self.yes {
                        Some(lv::Prompt::YesOrInteractive)
                    } else {
                        Some(lv::Prompt::Interactive)
                    }
                }
                _ => {
                    if self.yes {
                        Some(lv::Prompt::YesOrBlock)
                    } else {
                        Some(lv::Prompt::Block)
                    }
                }
            }
            #[cfg(not(feature = "prompt"))]
            {
                None
            }
        };
        super::print_init::init_options(self.color.unwrap_or_default(), level, prompt, log_config);
    }

    pub fn print_level(&self) -> lv::Print {
        let level = (self.verbose as i8 - self.quiet as i8).clamp(-2, 2);
        level.into()
    }

    /// Merge `other` into self. Options in other will be applied on top of self (equivalent
    /// to specifying `self` then specify `other`
    pub fn merge(&mut self, other: &Self) {
        self.verbose += other.verbose;
        self.quiet += other.quiet;
        if let Some(color) = other.color {
            self.color = Some(color);
        }
        #[cfg(feature = "prompt")]
        {
            if other.yes {
                self.yes = true;
            }
            self.non_interactive += other.non_interactive;
            self.interactive += other.interactive;
        }
    }
}

/// Entry point to CLI
///
/// # Safety
/// A safe wrapper is generated by the [`cu::cli`](macro@crate::cli) macro.
/// See [module level documentation](self) for more.
#[inline(always)]
#[doc(hidden)]
pub unsafe fn __run<
    TArg: clap::Parser,
    TLogConfig: LogConfig + Send + Sync + 'static,
    FPreproc: FnOnce(&mut TArg),
    FLogConfig: FnOnce(&Flags) -> TLogConfig,
    FExecute: FnOnce(TArg) -> crate::Result<()>,
    FFlag: FnOnce(&TArg) -> &Flags,
>(
    fn_preproc: FPreproc,
    fn_log_config: FLogConfig,
    fn_execute: FExecute,
    fn_flag: FFlag,
) -> std::process::ExitCode {
    let start = std::time::Instant::now();
    let args = unsafe {
        parse_args_or_help::<TArg, TLogConfig, FPreproc, FLogConfig, FFlag>(
            fn_preproc,
            fn_log_config,
            fn_flag,
        )
    };
    let result = fn_execute(args);
    handle_result(start, result)
}

/// Entry point to CLI
///
/// # Safety
/// A safe wrapper is generated by the [`cu::cli`](macro@crate::cli) macro.
/// See [module level documentation](self) for more.
#[inline(always)]
#[cfg(feature = "coroutine")]
#[doc(hidden)]
pub unsafe fn __co_run<
    TArg: clap::Parser + Send + 'static,
    TLogConfig: LogConfig + Send + Sync + 'static,
    FPreproc: FnOnce(&mut TArg),
    FLogConfig: FnOnce(&Flags) -> TLogConfig,
    FExecute: FnOnce(TArg) -> TResult + Send + 'static,
    TResult: Future<Output = crate::Result<()>> + Send + 'static,
    FFlag: FnOnce(&TArg) -> &Flags,
>(
    fn_preproc: FPreproc,
    fn_log_config: FLogConfig,
    fn_execute: FExecute,
    fn_flag: FFlag,
) -> std::process::ExitCode {
    let start = std::time::Instant::now();
    let args = unsafe {
        parse_args_or_help::<TArg, TLogConfig, FPreproc, FLogConfig, FFlag>(
            fn_preproc,
            fn_log_config,
            fn_flag,
        )
    };
    #[cfg(not(feature = "coroutine-heavy"))]
    let result = crate::co::block(async move { fn_execute(args).await });
    #[cfg(feature = "coroutine-heavy")]
    let result = crate::co::run(async move { fn_execute(args).await });

    handle_result(start, result)
}

unsafe fn parse_args_or_help<
    TArg: Parser,
    TLogConfig: LogConfig + Send + Sync + 'static,
    FPreproc: FnOnce(&mut TArg),
    FLogConfig: FnOnce(&Flags) -> TLogConfig,
    FFlag: FnOnce(&TArg) -> &Flags,
>(
    fn_preproc: FPreproc,
    fn_log_config: FLogConfig,
    fn_flag: FFlag,
) -> TArg {
    let mut parsed = parse_args::<TArg>();
    fn_preproc(&mut parsed);
    let flags = fn_flag(&parsed);
    let log_config: Arc<dyn LogConfig + Send + Sync> = Arc::new(fn_log_config(flags));
    unsafe { flags.apply(log_config) };
    parsed
}

/// Wrapper for clap parse to respect the color flag when printing help or error
fn parse_args<T: Parser>() -> T {
    // parse the color arg first, so that we can respect it when printing help
    let color = lv::Color::from_os_args();
    let use_color = color.is_colored_for_stdout();

    // this will exit on error
    let mut matches = get_colored_command::<T>(use_color).get_matches();

    match <T as FromArgMatches>::from_arg_matches_mut(&mut matches) {
        Ok(x) => x,
        Err(e) => {
            let mut command = get_colored_command::<T>(use_color);
            let error = e.format(&mut command);
            error.exit()
        }
    }
}

/// Try to parse arguments from an iterator and print the error/help
/// on failure.
///
/// Whether the output has color depends on the main CLI `--color` option.
/// This is useful for implementing custom command parser within
/// an application.
pub fn try_parse<T: Parser, I: IntoIterator>(iter: I) -> Option<T>
where
    I::Item: Into<OsString> + Clone,
{
    let use_color = crate::lv::color_enabled();
    let result = get_colored_command::<T>(use_color)
        .try_get_matches_from(iter)
        .and_then(|mut matches| <T as FromArgMatches>::from_arg_matches_mut(&mut matches));
    match result {
        Ok(x) => Some(x),
        Err(e) => {
            let mut command = get_colored_command::<T>(use_color);
            let error = e.format(&mut command);
            if let Err(e) = error.print() {
                crate::warn!("arg parse error failed to print: {e:?}");
            }
            None
        }
    }
}

/// Print the help text from a command parser.
///
/// Whether the output has color depends on the main CLI `--color` option.
/// This is useful for implementing custom command parser within
/// an application.
#[inline(always)]
pub fn print_help<T: Parser>(long: bool) {
    let command = get_colored_command::<T>(crate::lv::color_enabled());
    print_help_impl(command, long)
}
fn print_help_impl(mut command: Command, long: bool) {
    let result = if long {
        command.print_long_help()
    } else {
        command.print_help()
    };
    if let Err(e) = result {
        crate::warn!("help failed to print: {e:?}");
    }
}

#[inline(always)]
fn get_colored_command<T: Parser>(color: bool) -> Command {
    get_colored_command_impl(<T as CommandFactory>::command(), color)
}
fn get_colored_command_impl(command: Command, color: bool) -> Command {
    use clap::builder::styling::{AnsiColor, Styles};
    if color {
        // Modified version of Cargo's color style
        // [source](https://github.com/crate-ci/clap-cargo/blob/master/src/style.rs)
        command.styles(
            Styles::styled()
                .header(AnsiColor::BrightYellow.on_default())
                .usage(AnsiColor::BrightRed.on_default())
                .literal(AnsiColor::BrightCyan.on_default())
                .placeholder(AnsiColor::Cyan.on_default())
                .error(AnsiColor::BrightRed.on_default())
                .valid(AnsiColor::BrightCyan.on_default())
                .invalid(AnsiColor::BrightYellow.on_default())
                .context(AnsiColor::BrightYellow.on_default()),
        )
    } else {
        command.styles(Styles::plain())
    }
}

fn handle_result(start: Instant, result: crate::Result<()>) -> std::process::ExitCode {
    let elapsed = start.elapsed().as_secs_f32();
    if let Err(e) = result {
        crate::error!("fatal: {e:?}");
        // we display the hint for user to use -vv
        // if:
        // - the user is already tried to get more debug info with -v
        //   (because otherwise it will be too noisy and it might be a user-error,
        //   not a bug
        // - the trace hint is not explicitly disabled
        // - the trace hint is not already displayed
        if lv::D.enabled() && crate::lv::is_trace_hint_enabled() {
            if std::env::var("RUST_BACKTRACE")
                .unwrap_or_default()
                .is_empty()
            {
                crate::hint!("use -vv or set RUST_BACKTRACE=1 to display the error backtrace.");
            }
        }
        if crate::lv::is_print_time_enabled() {
            // use debug so the error trace is the last line,
            // so user is directed to see what is the most important
            crate::debug!("finished in {elapsed:.2}s");
        }
        reset_color();
        std::process::ExitCode::FAILURE
    } else {
        if crate::lv::is_print_time_enabled() {
            crate::info!("finished in {elapsed:.2}s");
        }
        reset_color();
        std::process::ExitCode::SUCCESS
    }
}

fn reset_color() {
    use std::io::IsTerminal as _;
    use std::io::Write as _;
    let mut stdout = std::io::stdout();
    if stdout.is_terminal() {
        let _ = write!(stdout, "\x1b[0m");
        let _ = stdout.flush();
        return;
    }
    let mut stderr = std::io::stderr();
    if stderr.is_terminal() {
        let _ = write!(stderr, "\x1b[0m");
        let _ = stderr.flush();
    }
}