fix logging

guidance-ai · Jul 31, 2024 · 9519650 · 9519650
1 parent 5e66bc7
commit 9519650
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 17 deletions.
diff --git a/parser/src/earley/from_guidance.rs b/parser/src/earley/from_guidance.rs
@@ -220,13 +220,20 @@ pub fn grammars_from_json(
         .enumerate()
         .map(|(idx, (lex, mut grm))| {
             if logger.level_enabled(2) {
-                writeln!(logger, "Grammar #{}:\n{:?}\n{:?}\n", idx, lex, grm).unwrap();
+                writeln!(
+                    logger.info_logger(),
+                    "Grammar #{}:\n{:?}\n{:?}\n",
+                    idx,
+                    lex,
+                    grm
+                )
+                .unwrap();
             }
 
             grm = grm.optimize();
 
             if logger.level_enabled(2) {
-                write!(logger, "  == Optimize ==>\n{:?}", grm).unwrap();
+                write!(logger.info_logger(), "  == Optimize ==>\n{:?}", grm).unwrap();
             }
 
             Arc::new(grm.compile(lex))

diff --git a/parser/src/lib.rs b/parser/src/lib.rs
@@ -13,7 +13,7 @@ macro_rules! infoln {
     ($s:expr, $($arg:tt)*) => {
         if $s.logger.level_enabled(2) {
             use std::fmt::Write;
-            writeln!($s.logger, $($arg)*).unwrap();
+            writeln!($s.logger.info_logger(), $($arg)*).unwrap();
         }
     };
 }
@@ -23,8 +23,8 @@ macro_rules! warn {
     ($s:expr, $($arg:tt)*) => {
         if $s.logger.level_enabled(1) {
             use std::fmt::Write;
-            $s.logger.write_str("Warning: ").unwrap();
-            writeln!($s.logger, $($arg)*).unwrap();
+            $s.logger.write_warning("Warning: ");
+            writeln!($s.logger.warning_logger(), $($arg)*).unwrap();
         }
     };
 }
diff --git a/parser/src/logging.rs b/parser/src/logging.rs
@@ -7,6 +7,14 @@ pub struct Logger {
     buffer: String,
 }
 
+pub struct InfoLogger<'a> {
+    logger: &'a mut Logger,
+}
+
+pub struct WarningLogger<'a> {
+    logger: &'a mut Logger,
+}
+
 impl Clone for Logger {
     fn clone(&self) -> Self {
         Self {
@@ -30,16 +38,38 @@ impl Logger {
 
     pub fn warn(&mut self, s: &str) {
         if self.level_enabled(1) {
-            self.write_str("Warning: ").unwrap();
-            self.write_str(s).unwrap();
-            self.write_str("\n").unwrap();
+            self.write_warning("Warning: ");
+            self.write_warning(s);
+            self.write_warning("\n");
         }
     }
 
     pub fn info(&mut self, s: &str) {
         if self.level_enabled(2) {
-            self.write_str(s).unwrap();
-            self.write_str("\n").unwrap();
+            self.write_info(s);
+            self.write_info("\n");
+        }
+    }
+
+    pub fn write_warning(&mut self, s: &str) {
+        self.write_level(1, s);
+    }
+
+    pub fn write_info(&mut self, s: &str) {
+        self.write_level(2, s);
+    }
+
+    pub fn write_buffer(&mut self, s: &str) {
+        self.buffer.push_str(s);
+    }
+
+    #[inline(always)]
+    pub fn write_level(&mut self, level: u32, s: &str) {
+        if self.buffer_level >= level {
+            self.buffer.push_str(s);
+        }
+        if self.stderr_level >= level {
+            eprint!("{}", s);
         }
     }
 
@@ -80,16 +110,26 @@ impl Logger {
     pub fn get_and_clear_logs(&mut self) -> String {
         std::mem::take(&mut self.buffer)
     }
+
+    pub fn info_logger(&mut self) -> InfoLogger {
+        InfoLogger { logger: self }
+    }
+
+    pub fn warning_logger(&mut self) -> WarningLogger {
+        WarningLogger { logger: self }
+    }
 }
 
-impl Write for Logger {
+impl<'a> Write for InfoLogger<'a> {
     fn write_str(&mut self, s: &str) -> std::fmt::Result {
-        if self.effective_level >= self.buffer_level {
-            self.buffer.push_str(s);
-        }
-        if self.effective_level >= self.stderr_level {
-            eprint!("{}", s);
-        }
+        self.logger.write_info(s);
+        Ok(())
+    }
+}
+
+impl<'a> Write for WarningLogger<'a> {
+    fn write_str(&mut self, s: &str) -> std::fmt::Result {
+        self.logger.write_warning(s);
         Ok(())
     }
 }