Improve logging config for pyo3

nautechsystems · Jan 16, 2024 · 1fa7ea2 · 1fa7ea2
1 parent 51517fe
commit 1fa7ea2
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 68 deletions.
diff --git a/nautilus_core/common/src/ffi/logging.rs b/nautilus_core/common/src/ffi/logging.rs
@@ -13,11 +13,7 @@
 //  limitations under the License.
 // -------------------------------------------------------------------------------------------------
 
-use std::{
-    collections::HashMap,
-    ffi::{c_char, CStr},
-    str::FromStr,
-};
+use std::ffi::{c_char, CStr};
 
 use log::LevelFilter;
 use nautilus_core::{
@@ -28,12 +24,12 @@ use nautilus_core::{
     uuid::UUID4,
 };
 use nautilus_model::identifiers::trader_id::TraderId;
-use serde_json::Value;
-use ustr::Ustr;
 
 use crate::{
     enums::{LogColor, LogLevel},
-    logging::{self, FileWriterConfig, LoggerConfig},
+    logging::{
+        self, map_log_level_to_filter, parse_component_levels, FileWriterConfig, LoggerConfig,
+    },
 };
 
 /// Returns whether the core logger is enabled.
@@ -116,41 +112,6 @@ pub unsafe extern "C" fn logging_init(
     logging::init_logging(trader_id, instance_id, config, file_config);
 }
 
-fn map_log_level_to_filter(log_level: LogLevel) -> LevelFilter {
-    match log_level {
-        LogLevel::Debug => LevelFilter::Debug,
-        LogLevel::Info => LevelFilter::Info,
-        LogLevel::Warning => LevelFilter::Warn,
-        LogLevel::Error => LevelFilter::Error,
-    }
-}
-
-fn parse_level_filter_str(s: &str) -> LevelFilter {
-    let mut log_level_str = s.to_string().to_uppercase();
-    if log_level_str == "WARNING" {
-        log_level_str = "WARN".to_string()
-    }
-    LevelFilter::from_str(&log_level_str)
-        .unwrap_or_else(|_| panic!("Invalid `LevelFilter` string, was {log_level_str}"))
-}
-
-fn parse_component_levels(
-    original_map: Option<HashMap<String, Value>>,
-) -> HashMap<Ustr, LevelFilter> {
-    match original_map {
-        Some(map) => {
-            let mut new_map = HashMap::new();
-            for (key, value) in map {
-                let ustr_key = Ustr::from(&key);
-                let value = parse_level_filter_str(&value.to_string());
-                new_map.insert(ustr_key, value);
-            }
-            new_map
-        }
-        None => HashMap::new(),
-    }
-}
-
 /// Creates a new log event.
 ///
 /// # Safety

diff --git a/nautilus_core/common/src/logging.rs b/nautilus_core/common/src/logging.rs
@@ -165,6 +165,41 @@ impl FileWriterConfig {
     }
 }
 
+pub fn map_log_level_to_filter(log_level: LogLevel) -> LevelFilter {
+    match log_level {
+        LogLevel::Debug => LevelFilter::Debug,
+        LogLevel::Info => LevelFilter::Info,
+        LogLevel::Warning => LevelFilter::Warn,
+        LogLevel::Error => LevelFilter::Error,
+    }
+}
+
+pub fn parse_level_filter_str(s: &str) -> LevelFilter {
+    let mut log_level_str = s.to_string().to_uppercase();
+    if log_level_str == "WARNING" {
+        log_level_str = "WARN".to_string()
+    }
+    LevelFilter::from_str(&log_level_str)
+        .unwrap_or_else(|_| panic!("Invalid `LevelFilter` string, was {log_level_str}"))
+}
+
+pub fn parse_component_levels(
+    original_map: Option<HashMap<String, serde_json::Value>>,
+) -> HashMap<Ustr, LevelFilter> {
+    match original_map {
+        Some(map) => {
+            let mut new_map = HashMap::new();
+            for (key, value) in map {
+                let ustr_key = Ustr::from(&key);
+                let value = parse_level_filter_str(value.as_str().unwrap());
+                new_map.insert(ustr_key, value);
+            }
+            new_map
+        }
+        None => HashMap::new(),
+    }
+}
+
 /// Initialize tracing.
 ///
 /// Tracing is meant to be used to trace/debug async Rust code. It can be

diff --git a/nautilus_core/common/src/python/logging.rs b/nautilus_core/common/src/python/logging.rs
@@ -13,15 +13,19 @@
 //  limitations under the License.
 // -------------------------------------------------------------------------------------------------
 
-use std::borrow::Cow;
+use std::{borrow::Cow, collections::HashMap};
 
-use nautilus_core::time::UnixNanos;
+use log::LevelFilter;
+use nautilus_core::{time::UnixNanos, uuid::UUID4};
+use nautilus_model::identifiers::trader_id::TraderId;
 use pyo3::prelude::*;
 use ustr::Ustr;
 
 use crate::{
     enums::{LogColor, LogLevel},
-    logging::{self, FileWriterConfig, LoggerConfig},
+    logging::{
+        self, map_log_level_to_filter, parse_level_filter_str, FileWriterConfig, LoggerConfig,
+    },
 };
 
 /// Initialize tracing.
@@ -51,27 +55,56 @@ pub fn py_init_tracing() {
 ///
 /// Should only be called once during an applications run, ideally at the
 /// beginning of the run.
-// #[pyfunction]
-// #[pyo3(name = "init_logging")]
-// pub fn py_init_logging(
-//     trader_id: TraderId,
-//     instance_id: UUID4,
-//     level_stdout: LogLevel,
-//     level_file: LogLevel,
-//     file_logging: u8,
-//     directory: Option<String>,
-//     file_name: Option<String>,
-//     file_format: Option<String>,
-// ) {
-//     logging::init_logging(
-//         trader_id,
-//         instance_id,
-//         config_spec,
-//         directory,
-//         file_name,
-//         file_format,
-//     );
-// }
+#[pyfunction]
+#[pyo3(name = "init_logging")]
+#[allow(clippy::too_many_arguments)]
+pub fn py_init_logging(
+    trader_id: TraderId,
+    instance_id: UUID4,
+    level_stdout: LogLevel,
+    level_file: Option<LogLevel>,
+    component_levels: Option<HashMap<String, String>>,
+    directory: Option<String>,
+    file_name: Option<String>,
+    file_format: Option<String>,
+    is_colored: Option<bool>,
+    is_bypassed: Option<bool>,
+    print_config: Option<bool>,
+) {
+    let level_file = level_file
+        .map(map_log_level_to_filter)
+        .unwrap_or(LevelFilter::Off);
+
+    let config = LoggerConfig::new(
+        map_log_level_to_filter(level_stdout),
+        level_file,
+        parse_component_levels(component_levels),
+        is_colored.unwrap_or(true),
+        is_bypassed.unwrap_or(false),
+        print_config.unwrap_or(false),
+    );
+
+    let file_config = FileWriterConfig::new(directory, file_name, file_format);
+
+    logging::init_logging(trader_id, instance_id, config, file_config);
+}
+
+fn parse_component_levels(
+    original_map: Option<HashMap<String, String>>,
+) -> HashMap<Ustr, LevelFilter> {
+    match original_map {
+        Some(map) => {
+            let mut new_map = HashMap::new();
+            for (key, value) in map {
+                let ustr_key = Ustr::from(&key);
+                let value = parse_level_filter_str(&value);
+                new_map.insert(ustr_key, value);
+            }
+            new_map
+        }
+        None => HashMap::new(),
+    }
+}
 
 /// Create a new log event.
 #[pyfunction]

diff --git a/nautilus_core/common/src/python/mod.rs b/nautilus_core/common/src/python/mod.rs
@@ -34,7 +34,7 @@ pub fn common(_: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_class::<LoggerConfig>()?;
     m.add_class::<FileWriterConfig>()?;
     m.add_function(wrap_pyfunction!(logging::py_init_tracing, m)?)?;
-    // m.add_function(wrap_pyfunction!(logging::py_init_logging, m)?)?;  // TODO: Will repair
+    m.add_function(wrap_pyfunction!(logging::py_init_logging, m)?)?;
     m.add_function(wrap_pyfunction!(logging::py_logger_log, m)?)?;
 
     Ok(())