Skip to main content

pistonite_cu/cli/
print_init.rs

1use std::sync::atomic::Ordering;
2use std::sync::{Arc, OnceLock};
3
4use cu::cli::printer::{PRINTER, Printer};
5#[cfg(feature = "prompt")]
6use cu::cli::prompt::PROMPT_LEVEL;
7use cu::lv;
8use env_filter::{Builder as LogEnvBuilder, Filter as LogEnvFilter};
9
10static LOGGER: OnceLock<LogImpl> = OnceLock::new();
11
12/// Shorthand to quickly setup logging. Can be useful in tests.
13///
14/// `qq`, `q`, `v` and `vv` inputs map to corresponding print levels. Other inputs
15/// are mapped to default level
16#[doc(alias = "quick_init")]
17pub fn level(lv: &str) {
18    let level = match lv {
19        "qq" => lv::Print::QuietQuiet,
20        "q" => lv::Print::Quiet,
21        "v" => lv::Print::Verbose,
22        "vv" => lv::Print::VerboseVerbose,
23        _ => lv::Print::Normal,
24    };
25    init_options(
26        lv::Color::Auto,
27        level,
28        Some(lv::Prompt::Block),
29        Arc::new(DefaultLogConfig),
30    );
31}
32
33/// Set global CLI options
34///
35/// **You do not need to call this if using the `#[cu::cli]` macro**.
36/// This is useful to initialize logging without using a cu-compatible command line args,
37/// say for a program that should not use clap to parse its arguments.
38///
39/// ```rust,no_run
40/// # use pistonite_cu as cu;
41/// use std::sync::Arc;
42///
43/// fn main() {
44///     cu::cli::init_options(
45///         cu::lv::Color::Auto,
46///         cu::lv::Print::Normal,
47///         None,
48///         Arc::new(cu::cli::DefaultLogConfig)
49///     );
50///
51///     cu::hint!("hello!");
52///
53///     cu::cli::reset_color();
54/// }
55/// ```
56///
57/// If prompt option is `None`, it will be `Interactive` unless env var `CI` is `true` or `1`, in which case it becomes `No`.
58/// Prompt option is ignored unless `prompt` feature is enabled
59///
60/// ## Difference to `#[cu::cli]` macro
61/// - The main function is not wrapped, meaning any `cu::Result` you got you need to display it
62///   yourself to the user.
63/// - Need to call [`cu::cli::reset_color`] manually to reset color.
64/// - RUST_BACKTRACE is not enabled automatically in trace level. To enable it you can
65///   check `lv::T::enabled()` after calling `init_options` and use `unsafe { std::env::set_var(...)
66///   }`
67/// - `finished in XX.XXs` time is not printed
68#[doc(alias = "init_cli")]
69pub fn init_options(
70    color: lv::Color,
71    level: lv::Print,
72    prompt: Option<lv::Prompt>,
73    log_config: Arc<dyn LogConfig + Send + Sync>,
74) {
75    // not using cu::env_var, since we are before log initialization
76    let env_rust_log = std::env::var("RUST_LOG");
77    let (log_level_filter, log_filter) = match env_rust_log {
78        Ok(value) if !value.is_empty() => {
79            let filter = LogEnvBuilder::new().parse(&value).build();
80            let log_level_filter = filter.filter();
81            (log_level_filter.max(level.into()), Some(filter))
82        }
83        _ => (level.into(), None),
84    };
85    log::set_max_level(log_level_filter);
86
87    let use_color = color.is_colored_for_stdout();
88    lv::USE_COLOR.store(use_color, Ordering::Release);
89    let printer = Printer::new(use_color);
90    if let Ok(mut g_printer) = PRINTER.lock() {
91        *g_printer = Some(printer);
92    }
93    #[cfg(feature = "prompt")]
94    {
95        let prompt = match prompt {
96            Some(x) => x,
97            None => {
98                let is_ci = std::env::var("CI")
99                    .map(|mut x| {
100                        x.make_ascii_lowercase();
101                        matches!(x.trim(), "true" | "1")
102                    })
103                    .unwrap_or_default();
104                if is_ci {
105                    lv::Prompt::Block
106                } else {
107                    lv::Prompt::Interactive
108                }
109            }
110        };
111        PROMPT_LEVEL.set(prompt)
112    }
113    #[cfg(not(feature = "prompt"))]
114    {
115        let _ = prompt;
116    }
117
118    lv::PRINT_LEVEL.set(level);
119
120    let _ = LOGGER.set(LogImpl {
121        filter: log_filter,
122        config: log_config,
123    });
124    log::set_logger(LOGGER.get().unwrap()).unwrap();
125}
126struct LogImpl {
127    filter: Option<LogEnvFilter>,
128    config: Arc<dyn LogConfig + Send + Sync>,
129}
130impl log::Log for LogImpl {
131    fn enabled(&self, metadata: &log::Metadata) -> bool {
132        match &self.filter {
133            Some(filter) => filter.enabled(metadata),
134            None => lv::Lv::from(metadata.level()).can_print(lv::PRINT_LEVEL.get()),
135        }
136    }
137
138    fn log(&self, record: &log::Record) {
139        let (level, show_module) = self.config.process(record);
140        if level != record.level().into() {
141            let metadata = log::Metadata::builder()
142                .level(level.into())
143                .target(record.metadata().target())
144                .build();
145            if !self.enabled(&metadata) {
146                return;
147            }
148        } else {
149            if !self.enabled(record.metadata()) {
150                return;
151            }
152        }
153        let message = if show_module {
154            // enable source location logging in trace messages
155            let mut message = String::new();
156            format_module_prefix(
157                &mut message,
158                record.module_path(),
159                record.file(),
160                record.line(),
161            );
162            use std::fmt::Write;
163            let _: Result<_, _> = write!(&mut message, "{}", record.args());
164            message
165        } else {
166            record.args().to_string()
167        };
168        if let Ok(mut printer) = PRINTER.lock() {
169            if let Some(printer) = printer.as_mut() {
170                printer.print_message(level, &message);
171            }
172        }
173    }
174
175    fn flush(&self) {}
176}
177
178fn format_module_prefix(
179    message: &mut String,
180    module: Option<&str>,
181    file: Option<&str>,
182    line: Option<u32>,
183) {
184    if module.is_none() && file.is_none() {
185        return;
186    }
187    message.push('[');
188    if let Some(p) = module {
189        // aliased crate, use the shorthand
190        if let Some(rest) = p.strip_prefix("pistonite_") {
191            message.push_str(rest);
192        } else {
193            message.push_str(p);
194        }
195        message.push(' ');
196    }
197    if let Some(f) = file {
198        let name = match f.rfind(['/', '\\']) {
199            None => f,
200            Some(i) => &f[i + 1..],
201        };
202        message.push_str(name);
203        if let Some(l) = line {
204            message.push(':');
205            message.push_str(&format!("{l}"));
206        }
207    }
208    message.push_str("] ");
209}
210
211/// Hook to configure the level and format before logging
212pub trait LogConfig {
213    /// Process a log record, return the level to log and if
214    /// the module path should be shown
215    fn process(&self, record: &lv::LogRecord) -> (lv::Lv, bool);
216}
217/// The default [`LogConfig`]
218pub struct DefaultLogConfig;
219impl LogConfig for DefaultLogConfig {
220    fn process(&self, record: &lv::LogRecord) -> (lv::Lv, bool) {
221        // downgrade logs from dependencies to trace
222        if let Some(m) = record.module_path() {
223            if m.starts_with("ignore") {
224                return (lv::T, true);
225            }
226        }
227
228        let level: lv::Lv = record.level().into();
229        (level, level == lv::T)
230    }
231}