xlog/lib/pretty.ml

open Core

type t = {
  mutex : Mutex.t;
  bp : Buffer.t;
  mutable align_to : int;
  out : out_channel;
  flags : int;
}

type config = {
  color : bool;
  timestamp : bool;
  namespace : bool;
  level : bool;
  source_loc : bool;
  backtrace : bool;
}

let _fl_c = 1  (* color *)
let _fl_t = 2  (* timestamp *)
let _fl_n = 4  (* namespace *)
let _fl_l = 8  (* level *)
let _fl_s = 16 (* source_loc *)
let _fl_b = 32 (* backtrace *)
let _fl_ct = _fl_c + _fl_t
let _fl_cl = _fl_c + _fl_l

let header = function
  | TRACE -> "TRACE"
  | DEBUG -> "DEBUG"
  | INFO  -> "INFO"
  | WARN  -> "WARN"
  | ERROR -> "ERROR"

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 box_at    = "└─ "
let box_exn   = "├─ "
let box_bt    = "├─── "

let pr_timestamp bp ts =
  Buffer.add_string bp (Time.to_string ts)

let pr_spaces bp n =
  for _ = 1 to n do
    Buffer.add_char bp ' '
  done

let rec pr_lines_rec indent bp s i =
  match String.index_from s i '\n' with
  | exception Not_found ->
    Buffer.add_substring bp s i (String.length s - i)
  | k ->
    Buffer.add_substring bp s i (k + 1 - i);
    pr_spaces bp indent;
    pr_lines_rec indent bp s (k + 1)

let pr_lines ~indent bp s =
  pr_lines_rec indent bp s 0

let make (out : out_channel) (cfg : config) = {
  mutex = Mutex.create ();
  bp = Buffer.create 512;
  align_to = 0;
  out;
  flags = (if cfg.color      then _fl_c else 0) +
          (if cfg.timestamp  then _fl_t else 0) +
          (if cfg.namespace  then _fl_n else 0) +
          (if cfg.level      then _fl_l else 0) +
          (if cfg.source_loc then _fl_s else 0) +
          (if cfg.backtrace  then _fl_b else 0)
}

let pr_msg bp msg ts ns lvl ~align ~indent ~f =
  begin
    if f _fl_ct then Buffer.add_string bp ansi_dim;
    if f _fl_t  then Printf.bprintf bp "%a " pr_timestamp ts;
    if f _fl_ct then Buffer.add_string bp ansi_off;

    if f _fl_n then Printf.bprintf bp "%s " ns;
    if f _fl_n then pr_spaces bp (align bp);

    if f _fl_cl then Buffer.add_string bp (ansi_header lvl);
    if f _fl_l then Printf.bprintf bp "%-5s " (header lvl);
    if f _fl_cl then Buffer.add_string bp ansi_off;

    if f _fl_c then Buffer.add_string bp ansi_bold;
    pr_lines bp msg ~indent:(indent bp);
    if f _fl_c then Buffer.add_string bp ansi_off;

    Buffer.add_string bp "\n";
  end

let writer t ~ts ~ns ~filename ~lineno ~func ~errno ~exn ~lvl msg =
  ignore errno;
  let f mask = t.flags land mask = mask in
  let align bp =
    let n = t.align_to - Buffer.length bp in
    t.align_to <- max t.align_to (Buffer.length bp);
    n
  in
  let indent bp =
    Buffer.length bp
    - (if f _fl_c  then 4 else 0)
    - (if f _fl_ct then 8 else 0)
    - (if f _fl_cl then 9 else 0)
  in
  begin
    Mutex.lock t.mutex;
    pr_msg t.bp msg ts ns lvl ~align ~indent ~f;
    begin match exn with
    | Some (exn, bt) ->
        Buffer.add_string t.bp box_exn;
        Printf.bprintf t.bp "Exception: %s\n" (Printexc.to_string exn);
        (* if not called within an exception handler, or backtraces are disabled, we don't get a
           backtrace and thus empty string will get added *)
        if f _fl_b then begin
          let bt_parts = String.split_on_char '\n' (Printexc.raw_backtrace_to_string bt) in
          List.iter (fun line ->
            if not (String.trim line = "") then begin
              Buffer.add_string t.bp box_bt;
              Buffer.add_string t.bp line;
              Buffer.add_string t.bp "\n" end) bt_parts
        end
    | None -> ()
    end;
    if f _fl_s && (not (filename = "") || not (func = "")) then begin
      Buffer.add_string t.bp box_at;
      begin match (filename, func) with
      | ("", _) -> Printf.bprintf t.bp "%s" func
      | (_, "") -> Printf.bprintf t.bp "%s:%d" filename lineno
      | _ -> Printf.bprintf t.bp "%s @ %s:%d" func filename lineno
      end;
      Buffer.add_string t.bp "\n"
    end;
    Buffer.output_buffer t.out t.bp;
    flush t.out;
    Buffer.clear t.bp;
    Mutex.unlock t.mutex
  end