pp: add Term_color extension

src/pp/containers_pp.ml

@@ -404,3 +404,83 @@ module Dump = struct
     let sep = char ';' ^ nl in
     group (char '[' ^ nest 1 (fill sep l) ^ char ']')
 end
+
+module Term_color = struct
+  type color =
+    [ `Black
+    | `Red
+    | `Yellow
+    | `Green
+    | `Blue
+    | `Magenta
+    | `Cyan
+    | `White
+    ]
+
+  type style =
+    [ `FG of color (* foreground *)
+    | `BG of color (* background *)
+    | `Bold
+    | `Reset
+    | `Underline
+    ]
+
+  open struct
+    let int_of_color_ = function
+      | `Black -> 0
+      | `Red -> 1
+      | `Green -> 2
+      | `Yellow -> 3
+      | `Blue -> 4
+      | `Magenta -> 5
+      | `Cyan -> 6
+      | `White -> 7
+
+    let code_of_style : style -> int = function
+      | `FG c -> 30 + int_of_color_ c
+      | `BG c -> 40 + int_of_color_ c
+      | `Bold -> 1
+      | `Reset -> 0
+      | `Underline -> 4
+
+    let spf = Printf.sprintf
+
+    let string_of_style a = spf "\x1b[%dm" (code_of_style a)
+
+    let reset = string_of_style `Reset
+
+    let string_of_style_list = function
+      | [] -> reset
+      | [ a ] -> string_of_style a
+      | [ a; b ] -> spf "\x1b[%d;%dm" (code_of_style a) (code_of_style b)
+      | [ a; b; c ] ->
+        spf "\x1b[%d;%d;%dm" (code_of_style a) (code_of_style b)
+          (code_of_style c)
+      | l ->
+        let buf = Buffer.create 32 in
+        let pp_num c =
+          Buffer.add_string buf (string_of_int (code_of_style c))
+        in
+        Buffer.add_string buf "\x1b[";
+        List.iteri
+          (fun i c ->
+            if i > 0 then Buffer.add_char buf ';';
+            pp_num c)
+          l;
+        Buffer.add_string buf "m";
+        Buffer.contents buf
+
+    (* TODO: handle nested styles *)
+    let ext_style_ : style list Ext.t =
+      {
+        pre = (fun out l -> Out.string out (string_of_style_list l));
+        post = (fun out _l -> Out.string out reset);
+      }
+  end
+
+  (** Set the foreground color. *)
+  let color (c : color) (d : t) : t = ext ext_style_ [ `FG c ] d
+
+  (** Set a full style for this document. *)
+  let style_l (l : style list) (d : t) : t = ext ext_style_ l d
+end

src/pp/containers_pp.mli

@@ -240,3 +240,29 @@ val sexp_l : t list -> t
 module Dump : sig
   val list : t list -> t
 end
+
+(** Simple colors in terminals *)
+module Term_color : sig
+  type color =
+    [ `Black
+    | `Blue
+    | `Cyan
+    | `Green
+    | `Magenta
+    | `Red
+    | `White
+    | `Yellow
+    ]
+
+  type style =
+    [ `BG of color
+    | `Bold
+    | `FG of color
+    | `Reset
+    | `Underline
+    ]
+
+  val color : color -> t -> t
+
+  val style_l : style list -> t -> t
+end