Source file `query.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
(** Terminal color querying utilities.

    Provides functions to query xterm-compatible terminals for their current
    foreground and background colors. *)

(** Terminal foreground and background colors. *)
type terminal_colours = {
  fg : (Gg.v4, string) result;  (** Foreground color or error message *)
  bg : (Gg.v4, string) result;  (** Background color or error message *)
}

(*
  xterm compatible terminals should support xlib queries
  which allow us to learn the current foreground and background
  colours used
  for more heuristics which may work on other terminals see:
  https://github.com/egberts/shell-term-background/blob/master/term-background.bash
*)

(** Xterm-compatible terminal query functions. *)
module Xterm = struct
  (** Set terminal to raw mode for direct input/output.
      @param set_when When to apply the settings (default: TCSAFLUSH)

      Translation of Python's `tty.setraw`
      https://github.com/python/cpython/blob/main/Lib/tty.py#L18 *)
  let set_raw ?(set_when=Unix.TCSAFLUSH) fd =
    let mode : Unix.terminal_io = {
      (Unix.tcgetattr fd) with
      c_brkint = false;
      c_icrnl = false;
      c_inpck = false;
      c_istrip = false;
      c_ixon = false;
      c_opost = false;
      c_csize = 8;
      c_parenb = false;
      c_echo = false;
      c_icanon = false;
      (* c_iexten = false; ...does not exist on Unix.terminal_io  *)
      c_ixoff = false; (* IEXTEN and IXOFF appear to set the same bit *)
      c_isig = false;
      c_vmin = 1;
      c_vtime = 0;
    } in
    Unix.tcsetattr fd set_when mode

  (** Query terminal using xterm control codes (internal implementation).
      Based on a Python version here:
      https://stackoverflow.com/a/45467190/202168 *)
  let query fd code =
    let fdname =
      if fd == Unix.stdin then "stdin"
      else if fd == Unix.stdout then "stdout"
      else if fd == Unix.stderr then "stderr"
      else "fd"
    in
    if Unix.isatty fd then
      let old_settings = Unix.tcgetattr fd in
      set_raw fd;
      Fun.protect
        ~finally:(fun () -> Unix.tcsetattr fd Unix.TCSADRAIN old_settings)
        (fun () ->
           Printf.printf "\o033]%s;?\o007" code;
           flush stdout;
           let r, _, _ = Unix.select [fd] [] [] 0.1 in
           let buf = Bytes.create 256 in
           (* Printf.printf ">> len r: %d\n" (List.length r); *)
           let readlen = match List.exists (fun (el) -> el == fd) r with
             | true -> Unix.read fd buf 0 256
             | false -> failwith @@ Printf.sprintf "Nothing to read on [%s]" fdname
           in
           (* Printf.printf ">> len buf: %d\n" readlen; *)
           Bytes.sub buf 0 readlen
           |> Bytes.escaped
           |> Bytes.to_string
        )
    else
      invalid_arg @@ Printf.sprintf "[%s] is not a tty" fdname

  (** Query terminal using xterm control codes.
      @param fd File descriptor to query (must be a TTY)
      @param code Control code to query (e.g., "10" for foreground, "11" for background)
      @return Result with terminal response string or error message *)
  let query fd code =
    try Ok ( query fd code )
    with Failure e | Invalid_argument e -> Error e

  (** Convert hexadecimal string to 8-bit integer (0-255) with scaling.
      Used for parsing xterm color responses.

      Translates a hexadecimal string, of any width, to an 8-bit int.
      The value will be 'scaled' according to number of hex chars,
      where each char is worth 4-bits.
      e.g.
        C -> C/F * FF = 204
        CCCC -> CCCC/FFFF * FF = 204
        C3 -> C3/FF * FF = 195
        C3B -> C3B/FFF * FF = 195
        C3C3 -> C3C3/FFFF * FF = 195
      see: https://stackoverflow.com/q/70962440/202168 *)
  let hex_to_8bit s =
    if String.length s = 0 then invalid_arg "hex_to_8bit: empty string";
    let scale = (16. ** float_of_int (String.length s)) -. 1. in
    let value = int_of_string @@ Printf.sprintf "0x%s" s in
    float_of_int value /. scale *. 255.
    |> Utils.int_round

  (** Parse xterm RGB color string (format: "rgb:RRRR/GGGG/BBBB").
      @param s Color string from terminal query response
      @return Result with parsed color or error message

      xterm returns colours in a 48-bit hex format *)
  let parse_colour s =
    let rex = Re.Pcre.re {|rgb:([0-9a-f]{1,4})/([0-9a-f]{1,4})/([0-9a-f]{1,4})|} |> Re.compile in
    match Re.exec_opt rex s with
    | Some groups ->
      let r = Re.Group.get groups 1 in
      let g = Re.Group.get groups 2 in
      let b = Re.Group.get groups 3 in
      Ok ( Color.Rgb.(v (hex_to_8bit r) (hex_to_8bit g) (hex_to_8bit b) |> to_gg) )
    | None -> Error (Printf.sprintf "Unrecognised colour string: %s" s)

  (** Get current terminal foreground and background colors.
      @param fd File descriptor to query (typically Unix.stdin)
      @return Record with fg and bg colors (or errors) *)
  let get_colours fd =
    {
      fg = query fd "10" |> Result.map parse_colour |> Result.join;
      bg = query fd "11" |> Result.map parse_colour |> Result.join;
    }
end