From 1712b51ecc5113e3762204a131ff78bb6b9cfd5a Mon Sep 17 00:00:00 2001
From: Tomer Shalev <tshalev@proofpoint.com>
Date: Wed, 7 Jul 2021 14:42:12 +0300
Subject: [PATCH] Add prefix customization to default formatter

---
 src/fmt/mod.rs | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs     |  9 ++++++-
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/src/fmt/mod.rs b/src/fmt/mod.rs
index 3c4fee0a..90a5c063 100644
--- a/src/fmt/mod.rs
+++ b/src/fmt/mod.rs
@@ -144,6 +144,7 @@ pub(crate) struct Builder {
     pub format_level: bool,
     pub format_indent: Option<usize>,
     pub custom_format: Option<FormatFn>,
+    pub format_prefix: Option<&'static str>,
     pub format_suffix: &'static str,
     built: bool,
 }
@@ -156,6 +157,7 @@ impl Default for Builder {
             format_level: true,
             format_indent: Some(4),
             custom_format: None,
+            format_prefix: None,
             format_suffix: "\n",
             built: false,
         }
@@ -189,6 +191,7 @@ impl Builder {
                     level: built.format_level,
                     written_header_value: false,
                     indent: built.format_indent,
+                    prefix: built.format_prefix,
                     suffix: built.format_suffix,
                     buf,
                 };
@@ -214,11 +217,13 @@ struct DefaultFormat<'a> {
     written_header_value: bool,
     indent: Option<usize>,
     buf: &'a mut Formatter,
+    prefix: Option<&'a str>,
     suffix: &'a str,
 }
 
 impl<'a> DefaultFormat<'a> {
     fn write(mut self, record: &Record) -> io::Result<()> {
+        self.write_prefix()?;
         self.write_timestamp()?;
         self.write_level(record)?;
         self.write_module_path(record)?;
@@ -257,6 +262,15 @@ impl<'a> DefaultFormat<'a> {
         }
     }
 
+    fn write_prefix(&mut self) -> io::Result<()> {
+        match self.prefix {
+            None => Ok(()),
+            Some(prefix) => {
+                self.write_header_value(prefix)
+            }
+        }
+    }
+
     fn write_level(&mut self, record: &Record) -> io::Result<()> {
         if !self.level {
             return Ok(());
@@ -412,6 +426,7 @@ mod tests {
             level: true,
             written_header_value: false,
             indent: None,
+            prefix: None,
             suffix: "\n",
             buf: &mut f,
         });
@@ -433,6 +448,7 @@ mod tests {
             level: false,
             written_header_value: false,
             indent: None,
+            prefix: None,
             suffix: "\n",
             buf: &mut f,
         });
@@ -454,6 +470,7 @@ mod tests {
             level: true,
             written_header_value: false,
             indent: Some(4),
+            prefix: None,
             suffix: "\n",
             buf: &mut f,
         });
@@ -475,6 +492,7 @@ mod tests {
             level: true,
             written_header_value: false,
             indent: Some(0),
+            prefix: None,
             suffix: "\n",
             buf: &mut f,
         });
@@ -496,6 +514,7 @@ mod tests {
             level: false,
             written_header_value: false,
             indent: Some(4),
+            prefix: None,
             suffix: "\n",
             buf: &mut f,
         });
@@ -503,6 +522,50 @@ mod tests {
         assert_eq!("log\n    message\n", written);
     }
 
+    #[test]
+    fn format_prefix() {
+        let writer = writer::Builder::new()
+            .write_style(WriteStyle::Never)
+            .build();
+
+        let mut f = Formatter::new(&writer);
+
+        let written = write(DefaultFormat {
+            timestamp: None,
+            module_path: false,
+            level: false,
+            written_header_value: false,
+            indent: None,
+            prefix: Some("PREFIX"),
+            suffix: "\n",
+            buf: &mut f,
+        });
+
+        assert_eq!("[PREFIX] log\nmessage\n", written);
+    }
+
+    #[test]
+    fn format_prefix_with_level_and_module_path() {
+        let writer = writer::Builder::new()
+            .write_style(WriteStyle::Never)
+            .build();
+
+        let mut f = Formatter::new(&writer);
+
+        let written = write(DefaultFormat {
+            timestamp: None,
+            module_path: true,
+            level: true,
+            written_header_value: false,
+            indent: None,
+            prefix: Some("PREFIX"),
+            suffix: "\n",
+            buf: &mut f,
+        });
+
+        assert_eq!("[PREFIX INFO  test::path] log\nmessage\n", written);
+    }
+
     #[test]
     fn format_suffix() {
         let writer = writer::Builder::new()
@@ -517,6 +580,7 @@ mod tests {
             level: false,
             written_header_value: false,
             indent: None,
+            prefix: None,
             suffix: "\n\n",
             buf: &mut f,
         });
@@ -538,6 +602,7 @@ mod tests {
             level: false,
             written_header_value: false,
             indent: Some(4),
+            prefix: None,
             suffix: "\n\n",
             buf: &mut f,
         });
diff --git a/src/lib.rs b/src/lib.rs
index 31ea7c3d..564dfcc1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -213,11 +213,12 @@
 //! ## Tweaking the default format
 //!
 //! Parts of the default format can be excluded from the log output using the [`Builder`].
-//! The following example excludes the timestamp from the log output:
+//! The following example excludes the timestamp from the log output, and sets a prefix:
 //!
 //! ```
 //! env_logger::builder()
 //!     .format_timestamp(None)
+//!     .format_prefix("PREFIX")
 //!     .init();
 //! ```
 //!
@@ -631,6 +632,12 @@ impl Builder {
         self.format_timestamp(Some(fmt::TimestampPrecision::Nanos))
     }
 
+    /// Configures the start of line prefix.
+    pub fn format_prefix(&mut self, prefix: &'static str) -> &mut Self {
+        self.format.format_prefix = Some(prefix);
+        self
+    }
+
     /// Configures the end of line suffix.
     pub fn format_suffix(&mut self, suffix: &'static str) -> &mut Self {
         self.format.format_suffix = suffix;