Source file console.ml

(** Color printing in terminals  *)
open Core

(* http://www.termsys.demon.co.uk/vtansi.htm *)
module Ansi = struct

  let kill_line () = print_string "\027[2K"

  let bell () = print_string "\007"
  let home_cursor () =   print_string "\027[0G"
  let save_cursor () = print_string "\027[s"
  let unsave_cursor () = print_string "\027[u"

  (* if it's good enough for git then it's good enough for us... *)
  let capable = lazy (Unix.isatty Unix.stdout &&
                      match Sys.getenv "TERM" with
                      | Some "dumb"
                      | None -> false
                      | Some _ -> true)

  module Attr = struct
    type color = [
      | `Black
      | `Red
      | `Green
      | `Yellow
      | `Blue
      | `Magenta
      | `Cyan
      | `White
    ]

    type attr = [
      | `Reset
      | `Bright
      | `Dim
      | `Underscore
      | `Blink
      | `Reverse
      | `Hidden
    ]
    type t = [
      | attr
      | color
      | `Bg of color
    ]

    let attr_to_int : attr -> int = function
      | `Reset      -> 0
      | `Bright     -> 1
      | `Dim        -> 2
      | `Underscore -> 4
      | `Blink      -> 5
      | `Reverse    -> 7
      | `Hidden     -> 8

    let fg_to_int : color -> int = function
      | `Black   -> 30
      | `Red     -> 31
      | `Green   -> 32
      | `Yellow  -> 33
      | `Blue    -> 34
      | `Magenta -> 35
      | `Cyan    -> 36
      | `White   -> 37

    let bg_to_int : color -> int = function
      | `Black   -> 40
      | `Red     -> 41
      | `Green   -> 42
      | `Yellow  -> 43
      | `Blue    -> 44
      | `Magenta -> 45
      | `Cyan    -> 46
      | `White   -> 47

    let to_int : t -> int = function
      | `Bg v -> bg_to_int v
      | #color as v -> fg_to_int v
      | #attr as v -> attr_to_int v

    let list_to_string : t list -> string = function
      | [] -> ""
      | l -> Printf.sprintf "\027[%sm"
               (String.concat ~sep:";"
                  (List.map l
                     ~f:(fun att -> string_of_int (to_int att))))
  end

  type color = Attr.color

  type attr = [
    | `Bright
    | `Dim
    | `Underscore
    | `Reverse
    | color
    | `Bg of color
  ]

  let string_with_attr style string =
    if style = [] then string
    else
      String.concat
        [ Attr.list_to_string (style :> Attr.t list)
        ; string
        ; Attr.list_to_string [`Reset]
        ]


  let output (style:attr list) oc s start len =
    if Lazy.force capable && style <> [] then begin
      Out_channel.output_string oc (Attr.list_to_string (style :> Attr.t list));
      Out_channel.output oc ~buf:s ~pos:start ~len;
      Out_channel.output_string oc (Attr.list_to_string [`Reset]);
      Out_channel.flush oc
    end else
      Out_channel.output oc ~buf:s ~pos:start ~len

  let output_string (style:attr list) oc s =
    if Lazy.force capable && style <> [] then begin
      Out_channel.output_string oc (Attr.list_to_string (style :> Attr.t list));
      Out_channel.output_string oc s;
      Out_channel.output_string oc (Attr.list_to_string [`Reset]);
      Out_channel.flush oc
    end else
      Out_channel.output_string oc s

  let fprintf (style:attr list) channel fmt =
    if Lazy.force capable && style <> [] then
      Printf.fprintf
        channel
        ( "%s" ^^ fmt ^^ "\027[0m%!")
        (Attr.list_to_string (style :> Attr.t list))
    else
      Printf.fprintf
        channel
        (fmt ^^ "%!")

  let eprintf style fmt = fprintf style stderr fmt
  let printf  style fmt = fprintf style stdout fmt

end

let is_color_tty () = Lazy.force Ansi.capable

module Columnize
    (In:sig
       type t
       val length : t -> int
     end) :
sig
  val iter : middle:(sep:In.t -> In.t -> int -> unit)
    -> last:(In.t -> int -> unit)
    -> sep:In.t
    -> In.t list
    -> int
    -> unit
end
=
struct
  let lines columns a = (Array.length a - 1) / columns + 1

  (** Size of an array printed out with this column configuration
      (lines*chars per column)
  *)
  let dim columns a =
    let lines = lines columns a in
    let rec loop cnt current acc =
      if cnt = Array.length a then
        List.rev (current::acc)
      else if cnt mod lines = 0 then
        loop (cnt+1) (In.length a.(cnt)) (current::acc)
      else
        loop (cnt+1) (max (In.length a.(cnt)) current) acc
    in
    lines,loop 1 (In.length a.(0)) []

  let rec line_len ~sep_len acc = function
    | []   -> acc
    | [v]  -> acc + v
    | h::t -> line_len ~sep_len (acc + sep_len + h) t

  let find_dim ~sep_len a max_len =
    let rec loop lines cols cnt =
      let (nlines,ncols) = dim (cnt+1) a in
      if nlines > lines || lines = 1 (** we are not gaining in vertical space
                                         anymore *)
         || line_len ~sep_len 0 ncols > max_len (** we are overflowing *)
      then
        Array.of_list cols
      else
        loop nlines ncols (cnt + 1)
    in
    let lines,cols = dim 1 a in
    loop lines cols 1

  let columnize a columns =
    let lines = lines columns a in
    let res = ref [] in
    for i = lines - 1 downto 0 do
      let line_acc = ref [] in
      for j = columns - 1 downto 0  do
        let pos =  i + j * lines in
        if pos < Array.length a then
          line_acc := a.(pos) :: !line_acc
      done;
      res := !line_acc :: !res
    done;
    !res

  let rec fold_line ~middle ~last sep acc padding line =
    match line,padding with
    | [v],len::_     ->
      last ~acc v (len - In.length v)
    | h::t,len::tlen ->
      fold_line ~middle ~last sep
        (middle ~acc ~sep h (len - In.length h)) tlen t
    | _              -> assert false

  let fold ~init ~middle ~last ~sep l max_len =
    if l = [] then
      init
    else
      let a = Array.of_list l in
      let columns = find_dim a ~sep_len:(In.length sep) max_len in
      let res = columnize a (Array.length columns) in
      List.fold_left res
        ~f:(fun acc line ->
          fold_line ~middle ~last sep acc (Array.to_list columns) line)
        ~init

  let iter ~middle ~last =
    fold ~init:() ~last:(fun ~acc:() -> last) ~middle:(fun ~acc:() -> middle)

end

let width () =
  match Linux_ext.get_terminal_size with
  | Result.Error _ -> `Not_available
  | Result.Ok _ when not (Unix.isatty Unix.stdout) -> `Not_a_tty
  | Result.Ok get_size -> `Cols (snd (get_size `Controlling))

let print_list oc l =
  match (width () :> [ `Cols of int | `Not_a_tty | `Not_available ]) with
  | `Not_a_tty | `Not_available ->
    List.iter l ~f:(fun (s,_) -> print_endline s)
  | `Cols cols ->
    let print_styled (s,style) =
      Ansi.output_string style oc s
    in
    let sep = "  ",[] in
    let last v _ = print_styled v; Out_channel.output_string oc "\n"
    and middle ~sep v pad_len =
      print_styled v;
      Out_channel.output_string oc (String.make pad_len ' ');
      print_styled sep
    in
    let module Col = Columnize
                       (struct
                         type t = string * Ansi.attr list
                         let length (s,_) = String.length s
                       end)
    in
    Col.iter ~sep ~last ~middle l cols