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)] fn 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 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}