From 5ae936e507295288f41f1856523cd2f997cad704 Mon Sep 17 00:00:00 2001
From: Yoshua Wuyts <yoshuawuyts@gmail.com>
Date: Fri, 19 Jul 2019 23:46:11 +0200
Subject: [PATCH] log key-value pairs

Signed-off-by: Yoshua Wuyts <yoshuawuyts@gmail.com>
---
 src/ndjson.rs | 29 +++++++++++++++++++++++++++--
 src/pretty.rs | 35 +++++++++++++++++++++++++++++------
 2 files changed, 56 insertions(+), 8 deletions(-)

diff --git a/src/ndjson.rs b/src/ndjson.rs
index 9d84510..db74342 100644
--- a/src/ndjson.rs
+++ b/src/ndjson.rs
@@ -1,6 +1,6 @@
 //! Print logs as ndjson.
 
-use log::{LevelFilter, Log, Metadata, Record};
+use log::{kv, LevelFilter, Log, Metadata, Record};
 use serde_json::Value;
 use std::collections::HashMap;
 use std::time;
@@ -49,7 +49,7 @@ impl Log for Logger {
 fn print_ndjson(record: &Record<'_>) {
     let msg = Msg {
         level: get_level(record.level()),
-        key_values: None,
+        key_values: format_kv_pairs(&record),
         time: time::UNIX_EPOCH.elapsed().unwrap().as_millis(),
         msg: record.args().to_string(),
     };
@@ -66,3 +66,28 @@ fn get_level(level: log::Level) -> u8 {
         Error => 50,
     }
 }
+
+fn format_kv_pairs(record: &Record) -> Option<HashMap<String, Value>> {
+    struct Visitor {
+        key_values: Option<HashMap<String, Value>>,
+    }
+
+    impl<'kvs> kv::Visitor<'kvs> for Visitor {
+        fn visit_pair(
+            &mut self,
+            key: kv::Key<'kvs>,
+            val: kv::Value<'kvs>,
+        ) -> Result<(), kv::Error> {
+            if let None = self.key_values {
+                self.key_values = Some(HashMap::new());
+            }
+            let kv = self.key_values.as_mut().unwrap();
+            kv.insert(key.to_string(), val.to_string().into());
+            Ok(())
+        }
+    }
+
+    let mut visitor = Visitor { key_values: None };
+    record.key_values().visit(&mut visitor).unwrap();
+    visitor.key_values
+}
diff --git a/src/pretty.rs b/src/pretty.rs
index 33aa985..ea5d8dc 100644
--- a/src/pretty.rs
+++ b/src/pretty.rs
@@ -1,7 +1,7 @@
 //! Pretty print logs.
 
 use console::style;
-use log::{Level, LevelFilter, Log, Metadata, Record};
+use log::{kv, Level, LevelFilter, Log, Metadata, Record};
 
 #[derive(Debug)]
 pub struct Logger {}
@@ -34,14 +34,37 @@ impl Log for Logger {
     fn flush(&self) {}
 }
 
-// TODO: use format_key_val to pretty print kv's from log crate
 fn pretty_print(record: &Record<'_>) {
-    println!("{}{}", format_message(&record), format_line(&record));
+    println!(
+        "{}{}{}",
+        format_message(&record),
+        format_line(&record),
+        format_kv_pairs(&record),
+    );
 }
 
-// TODO: use with key values that could eventually be put through log macros
-fn _format_key_val(key: &str, val: &str) -> String {
-    format!("   › {}: {}\n", style(key).magenta(), val)
+fn format_kv_pairs(record: &Record) -> String {
+    struct Visitor {
+        string: String,
+    }
+
+    impl<'kvs> kv::Visitor<'kvs> for Visitor {
+        fn visit_pair(
+            &mut self,
+            key: kv::Key<'kvs>,
+            val: kv::Value<'kvs>,
+        ) -> Result<(), kv::Error> {
+            let string = &format!("   › {}: {}\n", style(key).magenta(), val);
+            self.string.push_str(string);
+            Ok(())
+        }
+    }
+
+    let mut visitor = Visitor {
+        string: String::new(),
+    };
+    record.key_values().visit(&mut visitor).unwrap();
+    visitor.string
 }
 
 fn format_line(record: &Record<'_>) -> String {