use tags for color handling in CCFormat

src/core/CCFormat.ml

@@ -122,24 +122,8 @@ let to_string pp x =
   Format.pp_print_flush fmt ();
   Buffer.contents buf
 
-let sprintf format =
-  let buf = Buffer.create 64 in
-  let fmt = Format.formatter_of_buffer buf in
-  Format.kfprintf
-    (fun _fmt -> Format.pp_print_flush fmt (); Buffer.contents buf)
-    fmt
-    format
-
 let fprintf = Format.fprintf
 
-
-let ksprintf ~f fmt =
-  let buf = Buffer.create 32 in
-  let out = Format.formatter_of_buffer buf in
-  Format.kfprintf
-    (fun _ -> Format.pp_print_flush out (); f (Buffer.contents buf))
-    out fmt
-
 let stdout = Format.std_formatter
 let stderr = Format.err_formatter
 
@@ -181,12 +165,114 @@ let int_of_color_ = function
   | `Cyan -> 6
   | `White -> 7
 
-(* same as [pp], but in color [c] *)
+type style =
-let color_str c out s =
+  [ `FG of color (* foreground *)
-  let n = int_of_color_ c in
+  | `BG of color (* background *)
-  Format.fprintf out "\x1b[3%dm%s\x1b[0m" n s
+  | `Bold
+  | `Reset
+  ]
 
-(* same as [pp], but in bold color [c] *)
+let code_of_style : style -> int = function
-let bold_str c out s =
+  | `FG c -> 30 + int_of_color_ c
-  let n = int_of_color_ c in
+  | `BG c -> 40 + int_of_color_ c
-  Format.fprintf out "\x1b[3%d;1m%s\x1b[0m" n s
+  | `Bold -> 1
+  | `Reset -> 0
+
+let ansi_l_to_str_ = function
+  | [] -> "\x1b[0m"
+  | [a] -> Format.sprintf "\x1b[%dm" (code_of_style a)
+  | [a;b] -> Format.sprintf "\x1b[%d;%dm" (code_of_style a) (code_of_style b)
+  | l ->
+      let pp_num out c = int out (code_of_style c) in
+      to_string (list ~start:"\x1b[" ~stop:"m" ~sep:";" pp_num) l
+
+(* parse a tag *)
+let style_of_tag_ s = match String.trim s with
+  | "reset" -> [`Reset]
+  | "black" -> [`FG `Black]
+  | "red" -> [`FG `Red]
+  | "green" -> [`FG `Green]
+  | "yellow" -> [`FG `Yellow]
+  | "blue" -> [`FG `Blue]
+  | "magenta" -> [`FG `Magenta]
+  | "cyan" -> [`FG `Cyan]
+  | "white" -> [`FG `White]
+  | "Black" -> [`FG `Black]
+  | "Red" -> [`FG `Red; `Bold]
+  | "Green" -> [`FG `Green; `Bold]
+  | "Yellow" -> [`FG `Yellow; `Bold]
+  | "Blue" -> [`FG `Blue; `Bold]
+  | "Magenta" -> [`FG `Magenta; `Bold]
+  | "Cyan" -> [`FG `Cyan; `Bold]
+  | "White" -> [`FG `White; `Bold]
+  | s -> failwith ("unknown style: " ^ s)
+
+let color_enabled = ref false
+
+(* either prints the tag of [s] or delegate to [or_else] *)
+let mark_open_tag ~or_else s =
+  try
+    let style = style_of_tag_ s in
+    if !color_enabled then ansi_l_to_str_ style else ""
+  with Not_found -> or_else s
+
+let mark_close_tag ~or_else s =
+  try
+    let _ = style_of_tag_ s in (* check if it's indeed about color *)
+    if !color_enabled then ansi_l_to_str_ [`Reset] else ""
+  with Not_found -> or_else s
+
+(* add color handling to formatter [ppf] *)
+let set_color_tag_handling ppf =
+  let open Format in
+  let functions = pp_get_formatter_tag_functions ppf () in
+  let functions' = {functions with
+    mark_open_tag=(mark_open_tag ~or_else:functions.mark_open_tag);
+    mark_close_tag=(mark_close_tag ~or_else:functions.mark_close_tag);
+  } in
+  pp_set_mark_tags ppf true; (* enable tags *)
+  pp_set_formatter_tag_functions ppf functions'
+
+let set_color_default =
+  let first = ref true in
+  fun b ->
+    if b && not !color_enabled then (
+      color_enabled := true;
+      if !first then (
+        first := false;
+        set_color_tag_handling stdout;
+        set_color_tag_handling stderr;
+      );
+    ) else if not b && !color_enabled then color_enabled := false
+
+(*$R
+  set_color_default true;
+  let s = sprintf
+    "what is your @{<White>favorite color@}? @{<blue>blue@}! No, @{<red>red@}! Ahhhhhhh@."
+  in
+  assert_equal ~printer:CCFun.id
+    "what is your \027[37;1mfavorite color\027[0m? \027[34mblue\027[0m! No, \027[31mred\027[0m! Ahhhhhhh\n"
+    s
+*)
+
+let sprintf format =
+  let buf = Buffer.create 64 in
+  let fmt = Format.formatter_of_buffer buf in
+  if !color_enabled then set_color_tag_handling fmt;
+  Format.kfprintf
+    (fun _fmt -> Format.pp_print_flush fmt (); Buffer.contents buf)
+    fmt
+    format
+
+(*$T
+  sprintf "yolo %s %d" "a b" 42 = "yolo a b 42"
+  sprintf "%d " 0 = "0 "
+*)
+
+let ksprintf ~f fmt =
+  let buf = Buffer.create 32 in
+  let out = Format.formatter_of_buffer buf in
+  if !color_enabled then set_color_tag_handling out;
+  Format.kfprintf
+    (fun _ -> Format.pp_print_flush out (); f (Buffer.contents buf))
+    out fmt

src/core/CCFormat.mli

@@ -70,25 +70,50 @@ val map : ('a -> 'b) -> 'b printer -> 'a printer
 
   Use ANSI escape codes https://en.wikipedia.org/wiki/ANSI_escape_code
   to put some colors on the terminal.
-  We only allow styling of constant strings, because nesting is almost
-  impossible with ANSI codes (unless we maintain a stack of codes explicitely).
 
+  This uses {b tags} in format strings to specify  the style. Current styles
+  are the following:
+
+  {ul
+    {- "reset" resets style}
+    {- "black" }
+    {- "red" }
+    {- "green" }
+    {- "yellow" }
+    {- "blue" }
+    {- "magenta" }
+    {- "cyan" }
+    {- "white" }
+    {- "Black" bold black}
+    {- "Red" bold red }
+    {- "Green" bold green }
+    {- "Yellow" bold yellow }
+    {- "Blue" bold blue }
+    {- "Magenta" bold magenta }
+    {- "Cyan" bold cyan }
+    {- "White" bold white }
+  }
+
+  Example:
+
+  {[
+    set_color_default true;;
+
+    Format.printf
+      "what is your @{<White>favorite color@}? @{<blue>blue@}! No, @{<red>red@}! Ahhhhhhh@.";;
+   ]}
+
+  {b status: experimental}
   @since NEXT_RELEASE *)
 
-type color =
+val set_color_tag_handling : t -> unit
-  [ `Black
+(** adds functions to support color tags to the given formatter.
-  | `Red
+    @since NEXT_RELEASE *)
-  | `Yellow
-  | `Green
-  | `Blue
-  | `Magenta
-  | `Cyan
-  | `White
-  ]
 
-val color_str : color -> string printer
+val set_color_default : bool -> unit
-
+(** [set_color_default b] enables color handling on the standard formatters
-val bold_str : color -> string printer
+    (stdout, stderr) if [b = true] as well as on {!sprintf} formatters;
+    it disables the color handling if [b = false]. *)
 
 (** {2 IO} *)