Source file fmt.ml

(* (c) Daniel Bünzli *)

open Stdlib
open Fmt_meta

let invalid_arg fmt = Format.kasprintf invalid_arg fmt

let style_renderer_of_raw = function
  | "\x00" -> `None
  | "\x01" -> `Ansi
  | _ -> `None

let style_renderer_to_raw = function `None -> "\x00" | `Ansi -> "\x01"

let style_renderer ppf =
  let res = meta_raw (meta_store ppf) style_renderer_tag in
  style_renderer_of_raw res

let set_style_renderer ppf renderer =
  if ppf == Format.str_formatter
  then invalid_arg "Impossible to apply style on string formatter" ;

  let store = meta_store ppf in
  let style_renderer = style_renderer_to_raw renderer in

  set_meta ppf store ~style_renderer

let ansi_style_reset = "\x1b[m"

type standard =
  [ `Black | `Red | `Green | `Yellow | `Blue | `Magenta | `Cyan | `White ]

type bright = [ `Bright of standard ]

type bit8 = [ `bit8 of int * int * int ]

type bit24 = [ `bit24 of int * int * int ]

type grayscale = [ `Grayscale of int ]

type style =
  [ `None
  | `Style of [ `Fg | `Bg ] * [ standard | bright | bit8 | bit24 | grayscale ]
  ]

type rest = [ standard | bright | bit8 | grayscale ]

let ansi_style_code = function
  | `None -> ansi_style_reset
  | `Style (where, (#bit24 as color)) ->
      let (`bit24 (r, g, b)) = color in
      if r >= 0 && r <= 255 && g >= 0 && g <= 255 && b >= 0 && b <= 255
      then
        let where = match where with `Fg -> 38 | `Bg -> 48 in
        Format.asprintf "\x1b[%d;2;%d;%d;%dm" where r g b
      else invalid_arg "Invalid color: bit24(%d, %d, %d)" r g b
  | `Style (where, (#rest as color)) ->
      let where = match where with `Fg -> 38 | `Bg -> 48 in
      let color =
        match color with
        | `Black -> 0
        | `Red -> 1
        | `Green -> 2
        | `Yellow -> 3
        | `Blue -> 4
        | `Magenta -> 5
        | `Cyan -> 6
        | `White -> 7
        | `Bright color -> (
            match color with
            | `Black -> 8
            | `Red -> 9
            | `Green -> 10
            | `Yellow -> 11
            | `Blue -> 12
            | `Magenta -> 13
            | `Cyan -> 14
            | `White -> 15)
        | `bit8 (r, g, b) ->
            if r >= 0 && r <= 5 && g >= 0 && g <= 5 && b >= 0 && b <= 5
            then 16 + (36 * r) + (6 * g) + b
            else invalid_arg "Invalid color: bit8(%d, %d, %d)" r g b
        | `Grayscale n ->
            if n >= 0 && n <= 24
            then 232 + n
            else invalid_arg "Invalid color: Grayscale(%d)" n in
      Format.asprintf "\x1b[%d;5;%dm" where color

let styled style pp ppf =
  match style_renderer ppf with
  | `None -> Format.fprintf ppf "%a" pp
  | `Ansi ->
      let reset ppf = Format.fprintf ppf "@<0>%s" ansi_style_reset in
      Format.kfprintf reset ppf "@<0>%s%a" (ansi_style_code style) pp

let with_buffer ?like buf =
  let ppf = Format.formatter_of_buffer buf in
  match like with
  | None -> ppf
  | Some like ->
      set_meta_store ppf (meta_store like) ;
      ppf

let strf_like ppf fmt =
  let buf = Buffer.create 80 in
  let bppf = with_buffer ~like:ppf buf in
  let flush ppf =
    Format.pp_print_flush ppf () ;
    let s = Buffer.contents buf in
    Buffer.reset buf ;
    s in
  Format.kfprintf flush bppf fmt