From 0711561c0d8a54eca5e7465b32221c55acbfa8d6 Mon Sep 17 00:00:00 2001
From: tali <iitalics@gmail.com>
Date: Wed, 31 Jan 2024 12:28:47 -0500
Subject: [PATCH] refactor ansi colors in pretty logger

---
 lib/logging/logging.ml | 67 +++++++++++++++++++++++++++++++-----------
 1 file changed, 50 insertions(+), 17 deletions(-)

diff --git a/lib/logging/logging.ml b/lib/logging/logging.ml
index ad25a28..77cb366 100644
--- a/lib/logging/logging.ml
+++ b/lib/logging/logging.ml
@@ -113,9 +113,11 @@ module Pretty = struct
     out : out_channel;
   }
 
-  (* TODO: config timestamp *)
-  (* TODO: config color *)
-  (* TODO: config namespace *)
+  (* TODO: config these per writer *)
+  let _color = true
+  let _timestamp = true
+  let _namespace = true
+  let _level = true
 
   let header = function
     | TRACE -> "TRACE"
@@ -124,12 +126,16 @@ module Pretty = struct
     | WARN  -> "WARN"
     | ERROR -> "ERROR"
 
-  let ansi_color = function
-    | TRACE -> 34
-    | DEBUG -> 36
-    | INFO  -> 32
-    | WARN  -> 33
-    | ERROR -> 31
+  let ansi_header = function
+    | TRACE -> "\x1b[34m"
+    | DEBUG -> "\x1b[36m"
+    | INFO  -> "\x1b[32m"
+    | WARN  -> "\x1b[33m"
+    | ERROR -> "\x1b[31m"
+
+  let ansi_dim = "\x1b[2m"
+  let ansi_bold = "\x1b[1m"
+  let ansi_off = "\x1b[0m"
 
   let pr_timestamp bp ts =
     let ts_ms = int_of_float (ts *. 1000.0) mod 1000 in
@@ -162,17 +168,44 @@ module Pretty = struct
     out;
   }
 
+  let pr_msg bp msg ts ns lvl ~align ~indent =
+    begin
+      if _timestamp && _color then Buffer.add_string bp ansi_dim;
+      if _timestamp then Printf.bprintf bp "%a " pr_timestamp ts;
+      if _timestamp && _color then Buffer.add_string bp ansi_off;
+
+      if _namespace then Printf.bprintf bp "%s " ns;
+      if _namespace then pr_spaces bp (align bp);
+
+      if _level && _color then Buffer.add_string bp (ansi_header lvl);
+      if _level then Printf.bprintf bp "%-5s " (header lvl);
+      if _level && _color then Buffer.add_string bp ansi_off;
+
+      if _color then Buffer.add_string bp ansi_bold;
+      pr_lines bp msg ~indent:(indent bp);
+      if _color then Buffer.add_string bp ansi_off;
+
+      Buffer.add_string bp "\n";
+    end
+
   let writer t ~ts ~ns ~lvl msg =
+    let align bp =
+      let n = t.align_to - Buffer.length bp in
+      t.align_to <- Buffer.length bp;
+      n
+    in
+    let indent bp =
+      let subtract =
+        if _color then
+          4 + (if _timestamp then 8 else 0) + (if _level then 9 else 0)
+        else
+          0
+      in
+      Buffer.length bp - subtract
+    in
     begin
       Mutex.lock t.mutex;
-      Printf.bprintf t.bp "\x1b[1m%a\x1b[22m %s "
-        pr_timestamp ts ns;
-      pr_spaces t.bp (t.align_to - Buffer.length t.bp);
-      t.align_to <- Buffer.length t.bp;
-      Printf.bprintf t.bp "\x1b[%dm%-5s\x1b[39;1m "
-        (ansi_color lvl) (header lvl);
-      pr_lines t.bp msg ~indent:(t.align_to - 3);
-      Printf.bprintf t.bp "\x1b[0m\n";
+      pr_msg t.bp msg ts ns lvl ~align ~indent;
       Buffer.output_buffer t.out t.bp;
       flush t.out;
       Buffer.clear t.bp;