rusty_logging/
logger.rs

1use crate::error::LoggingError;
2use dynfmt::{Format, SimpleCurlyFormat};
3use pyo3::prelude::*;
4use pyo3::types::PyTuple;
5use pyo3::types::PyTupleMethods;
6use serde::{Deserialize, Serialize};
7use std::io;
8use std::str::FromStr;
9use tracing_subscriber;
10use tracing_subscriber::fmt::time::UtcTime;
11
12#[pyclass(eq)]
13#[derive(Debug, PartialEq, Clone, Serialize, Deserialize, Default)]
14pub enum LogLevel {
15    Debug,
16    #[default]
17    Info,
18    Warn,
19    Error,
20    Trace,
21}
22
23impl FromStr for LogLevel {
24    type Err = LoggingError;
25
26    fn from_str(s: &str) -> Result<Self, Self::Err> {
27        match s.to_lowercase().as_str() {
28            "debug" => Ok(LogLevel::Debug),
29            "info" => Ok(LogLevel::Info),
30            "warn" => Ok(LogLevel::Warn),
31            "error" => Ok(LogLevel::Error),
32            "trace" => Ok(LogLevel::Trace),
33            _ => Ok(LogLevel::Info),
34        }
35    }
36}
37
38#[pyclass(eq)]
39#[derive(Debug, PartialEq, Clone, Serialize, Deserialize, Default)]
40pub enum WriteLevel {
41    #[default]
42    Stdout,
43    Stderror,
44}
45
46#[allow(clippy::len_zero)] // len tends to be faster than is_empty in tests
47fn format_string(message: &str, args: &Vec<String>) -> String {
48    if args.len() > 0 {
49        SimpleCurlyFormat
50            .format(message, args)
51            .unwrap_or_else(|_| message.into())
52            .to_string()
53    } else {
54        message.to_string()
55    }
56}
57
58pub fn parse_args(args: &Bound<'_, PyTuple>) -> Option<Vec<String>> {
59    if args.is_empty() {
60        None
61    } else {
62        Some(args.iter().map(|x| x.to_string()).collect())
63    }
64}
65
66const DEFAULT_TIME_PATTERN: &str =
67    "[year]-[month]-[day]T[hour repr:24]:[minute]:[second]::[subsecond digits:4]";
68
69fn build_json_subscriber(
70    log_level: tracing::Level,
71    config: &LoggingConfig,
72) -> Result<(), LoggingError> {
73    let sub = tracing_subscriber::fmt()
74        .with_max_level(log_level)
75        .json()
76        .with_target(false)
77        .flatten_event(true)
78        .with_thread_ids(config.show_threads)
79        .with_timer(config.time_format()?);
80
81    if config.write_level == WriteLevel::Stderror {
82        sub.with_writer(io::stderr).try_init().map_err(|e| {
83            LoggingError::Error(format!("Failed to setup logging with error: {}", e))
84        })?;
85    } else {
86        sub.with_writer(io::stdout).try_init().map_err(|e| {
87            LoggingError::Error(format!("Failed to setup logging with error: {}", e))
88        })?;
89    }
90    Ok(())
91}
92
93fn build_subscriber(log_level: tracing::Level, config: &LoggingConfig) -> Result<(), LoggingError> {
94    let sub = tracing_subscriber::fmt()
95        .with_max_level(log_level)
96        .with_target(false)
97        .with_thread_ids(config.show_threads)
98        .with_timer(config.time_format()?);
99
100    if config.write_level == WriteLevel::Stderror {
101        sub.with_writer(io::stderr).try_init().map_err(|e| {
102            LoggingError::Error(format!("Failed to setup logging with error: {}", e))
103        })?;
104    } else {
105        sub.with_writer(io::stdout).try_init().map_err(|e| {
106            LoggingError::Error(format!("Failed to setup logging with error: {}", e))
107        })?;
108    }
109    Ok(())
110}
111
112pub fn setup_logging(config: &LoggingConfig) -> Result<(), LoggingError> {
113    let display_level = match &config.log_level {
114        LogLevel::Debug => tracing::Level::DEBUG,
115        LogLevel::Info => tracing::Level::INFO,
116        LogLevel::Warn => tracing::Level::WARN,
117        LogLevel::Error => tracing::Level::ERROR,
118        LogLevel::Trace => tracing::Level::TRACE,
119    };
120
121    if config.use_json {
122        build_json_subscriber(display_level, config)
123    } else {
124        build_subscriber(display_level, config)
125    }
126}
127
128#[pyclass]
129#[derive(Debug, Clone, Deserialize, Serialize)]
130pub struct LoggingConfig {
131    #[pyo3(get, set)]
132    show_threads: bool,
133
134    #[pyo3(get, set)]
135    log_level: LogLevel,
136
137    #[pyo3(get, set)]
138    write_level: WriteLevel,
139
140    #[pyo3(get, set)]
141    use_json: bool,
142}
143
144#[pymethods]
145impl LoggingConfig {
146    #[new]
147    #[pyo3(signature = (show_threads=None, log_level=None, write_level=WriteLevel::Stdout, use_json=false))]
148    pub fn new(
149        show_threads: Option<bool>,
150        log_level: Option<LogLevel>,
151        write_level: Option<WriteLevel>,
152        use_json: Option<bool>,
153    ) -> Self {
154        let show_threads = show_threads.unwrap_or(true);
155
156        let write_level = write_level.unwrap_or(WriteLevel::Stdout);
157        let use_json = use_json.unwrap_or(false);
158
159        let log_level = log_level.unwrap_or(std::env::var("LOG_LEVEL").map_or_else(
160            |_| LogLevel::Info,
161            |x| LogLevel::from_str(&x).unwrap_or(LogLevel::Info),
162        ));
163
164        LoggingConfig {
165            show_threads,
166            log_level,
167            write_level,
168            use_json,
169        }
170    }
171
172    #[staticmethod]
173    pub fn json_default() -> Self {
174        LoggingConfig::new(Some(true), None, None, Some(true))
175    }
176
177    #[staticmethod]
178    #[allow(clippy::should_implement_trait)]
179    pub fn default() -> Self {
180        LoggingConfig::new(Some(true), None, None, None)
181    }
182}
183
184impl LoggingConfig {
185    /// Create a new LoggingConfig with the given parameters when using within Rust
186    ///
187    /// # Arguments
188    ///
189    /// * `show_threads` - Whether to show thread ids in logs
190    /// * `log_level` - The log level to use
191    /// * `write_level` - The write level to use
192    /// * `use_json` - Whether to use json format for logs
193    ///
194    pub fn rust_new(
195        show_threads: bool,
196        log_level: LogLevel,
197        write_level: WriteLevel,
198        use_json: bool,
199    ) -> Self {
200        LoggingConfig {
201            show_threads,
202            log_level,
203            write_level,
204            use_json,
205        }
206    }
207    fn time_format(
208        &self,
209    ) -> Result<UtcTime<Vec<time::format_description::FormatItem<'static>>>, LoggingError> {
210        let formatter = UtcTime::new(
211            time::format_description::parse(DEFAULT_TIME_PATTERN).map_err(|e| {
212                LoggingError::Error(format!(
213                    "Failed to parse time format: {} with error: {}",
214                    DEFAULT_TIME_PATTERN, e
215                ))
216            })?,
217        );
218
219        Ok(formatter)
220    }
221}
222
223#[pyclass]
224pub struct RustyLogger {}
225
226#[pymethods]
227impl RustyLogger {
228    #[staticmethod]
229    #[pyo3(signature = (config=None))]
230    pub fn setup_logging(config: Option<LoggingConfig>) -> Result<(), LoggingError> {
231        let config = config.unwrap_or(LoggingConfig::default());
232        let _ = setup_logging(&config).is_ok();
233
234        Ok(())
235    }
236
237    #[staticmethod]
238    #[pyo3(signature = (config=None))]
239    pub fn get_logger(config: Option<LoggingConfig>) -> Result<Self, LoggingError> {
240        let config = config.unwrap_or(LoggingConfig::default());
241        let _ = setup_logging(&config).is_ok();
242
243        Ok(RustyLogger {})
244    }
245
246    #[pyo3(signature = (message, *args))]
247    pub fn info(&self, message: &str, args: &Bound<'_, PyTuple>) {
248        let args = parse_args(args);
249        let msg = match args {
250            Some(val) => format_string(message, &val),
251            None => message.to_string(),
252        };
253        tracing::info!(msg);
254    }
255
256    #[pyo3(signature = (message, *args))]
257    pub fn debug(&self, message: &str, args: &Bound<'_, PyTuple>) {
258        let args = parse_args(args);
259        let msg = match args {
260            Some(val) => format_string(message, &val),
261            None => message.to_string(),
262        };
263        tracing::debug!(msg);
264    }
265
266    #[pyo3(signature = (message, *args))]
267    pub fn warn(&self, message: &str, args: &Bound<'_, PyTuple>) {
268        let args = parse_args(args);
269        let msg = match args {
270            Some(val) => format_string(message, &val),
271            None => message.to_string(),
272        };
273        tracing::warn!(msg);
274    }
275
276    #[pyo3(signature = (message, *args))]
277    pub fn error(&self, message: &str, args: &Bound<'_, PyTuple>) {
278        let args = parse_args(args);
279        let msg = match args {
280            Some(val) => format_string(message, &val),
281            None => message.to_string(),
282        };
283        tracing::error!(msg);
284    }
285
286    #[pyo3(signature = (message, *args))]
287    pub fn trace(&self, message: &str, args: &Bound<'_, PyTuple>) {
288        let args = parse_args(args);
289        let msg = match args {
290            Some(val) => format_string(message, &val),
291            None => message.to_string(),
292        };
293        tracing::trace!(msg);
294    }
295}
296
297impl RustyLogger {
298    pub fn setup(config: Option<&LoggingConfig>) -> Result<(), LoggingError> {
299        let default_config = LoggingConfig::default();
300        let config = config.unwrap_or(&default_config);
301        let _ = setup_logging(&config).is_ok();
302
303        Ok(())
304    }
305}