Automatic generation of zsh completion from an extended Arg description

ocompl.ml

(* thx example for tagpix.pl *)

open Printf

let trim s =
  let len = String.length s in
  let rec aux = function
    | n when n = len -> ""
    | n -> 
        if s.[n] <> ' ' 
        then String.sub s n (len-n)
        else aux (n+1)
  in aux 0
    
(*Filename.generic_quote*)
let generic_quote whatquote quotequote s =
  let l = String.length s in
  let b = Buffer.create (l + 20) in
  for i = 0 to l - 1 do
    if s.[i] = whatquote
    then Buffer.add_string b quotequote
    else Buffer.add_char b  s.[i]
  done;
  Buffer.contents b

let quote_help s = List.fold_left (fun s (c,s2) -> generic_quote c s2 s)
  s ['\'',"\\'";']',"\\]";'"',"\\\""] 

let quote_colon s = generic_quote ':' "\\:" s

module Completion =
struct
  type compl_spec = 
    | CFile of string option
    | CDir of string option
    | CInt
    | CFloat
    | CString
    | CSymbol of string list
    | CTuple of (compl_spec * string) list

  type compl = 
      { compl : (compl_spec * string) option;
        options : string list;
        excluded : string list;
        help : string;
      }
    

  type prog_descr = 
      { name : string;
        infos : (string * string) list;
        rest : compl_spec * string;
        descs : compl list;
      }

  module Zsh =
    struct
      let print_compl ch (c,mesg) = 
        fprintf ch ":%s:" (quote_colon (quote_help mesg));
        let file dir pattern =
          let dir = if dir then "-/ " else "" in
          let pattern = match pattern with
            | None -> ""
            | Some s -> sprintf "-g '%s'" (quote_help s) in
          fprintf ch "_files %s %s" dir pattern in
        let anything () = fprintf ch ""(*"( )"*) in
        let choice l =
          fprintf ch "(%s)" (String.concat " " l) in
        (match c with
           | CFile s -> file false s
           | CDir s -> file true s
           | CInt -> anything ()
           | CFloat -> anything ()
           | CString -> anything ()
           | CSymbol l -> choice l
           | CTuple _ -> assert false)

      let print_not_empty_list form delim ch = function
        | [] -> ()
        | l -> fprintf ch form (String.concat delim l)

      let print_not_sing_list form delim ch = function
        | [] -> assert false
        | [a] -> fprintf ch "%s" a
        | l -> fprintf ch form (String.concat delim l)
            
            

      let print_rule ch c =
        fprintf ch "     %a'%a'\"[%s]"
          (print_not_empty_list "\"(%s)\"" " ") c.excluded
          (print_not_sing_list "{%s}" ",") c.options
          (quote_help (trim c.help));
        (match c.compl with
           | None -> ()
           | Some cmessg ->  print_compl ch cmessg);
        fprintf ch "\"\\\n"

      let write prog ch =
        fprintf ch
          "#compdef %s
##
## zsh completion for %s
## generated by ocompl

##
##
##
##

local curcontext state line cmds ret=1

##
_arguments -s -S \\\n" prog.name prog.name;
        fprintf ch
"     '(- *)'{-h,--help}'[Brief help message]' \\\n";
        List.iter (print_rule ch) prog.descs;
        fprintf ch 
"     \"*%a\" \\\n" print_compl prog.rest;
        fprintf ch
"     && ret=0

return ret

_%s \"$@\"
" prog.name
    end

end

let env_compl_mode = "COMPL_MODE"

let is_mode_compl =
  try
    let r =Sys.getenv env_compl_mode in
    r="yes"
  with Not_found -> false

module Arg_compl =
struct
  open Arg

  type spec =
    | Unit of (unit -> unit)     (* Call the function with unit argument *)
    | Bool of string * (bool -> unit)     (* Call the function with a bool argument *)
    | Set of bool ref            (* Set the reference to true *)
    | Clear of bool ref          (* Set the reference to false *)
    | String of string * (string -> unit) (* Call the function with a string argument *)
    | Set_string of string * string ref   (* Set the reference to the string argument *)
    | File of string * string option * (string -> unit)   (* Call the function with a string argument *)
    | Set_file of string * string option * string ref     (* Set the reference to the string argument *)
    | Dir of string * string option * (string -> unit)   (* Call the function with a string argument *)
    | Set_dir of string * string option * string ref     (* Set the reference to the string argument *)
    | Int of string * (int -> unit)       (* Call the function with an int argument *)
    | Set_int of string * int ref         (* Set the reference to the int argument *)
    | Float of string * (float -> unit)   (* Call the function with a float argument *)
    | Set_float of string * float ref     (* Set the reference to the float argument *)
    | Tuple of string * spec list         (* Take several arguments according to the
                                    spec list *)
    | Symbol of string * string list * (string -> unit)
        (* Take one of the symbols as argument and
           call the function with the symbol. *)
    | Rest of (string -> unit)   (* Stop interpreting keywords and call the
                                    function with each remaining argument *)
        
  type options = key * key list * spec * doc

 
    (** Common part*)


  let rec legacy_of_spec = function
    | Unit f -> Arg.Unit f  | Bool (_,f) -> Arg.Bool f 
    | Set f -> Arg.Set f | Clear f -> Arg.Clear f
    | Int (_,f) -> Arg.Int f | Set_int (_,f) -> Arg.Set_int f
    | Float (_,f) -> Arg.Float f | Set_float (_,f) -> Arg.Set_float f
    | Tuple (_,l) -> Arg.Tuple (List.map legacy_of_spec l) 
    | Symbol (_,l,f) -> Arg.Symbol (l,f)
    | Rest f -> Arg.Rest f
    | String (_,f) | File (_,_,f) | Dir (_,_,f) -> Arg.String f
    | Set_string (_,f) | Set_file (_,_,f) | Set_dir (_,_,f) -> Arg.Set_string f

  let rec spec_of_legacy_aux mesg = function
    | Arg.Unit f -> Unit f  | Arg.Bool f -> Bool (mesg,f)
    | Arg.Set f -> Set f | Arg.Clear f -> Clear f
    | Arg.Int f -> Int (mesg,f) | Arg.Set_int f -> Set_int (mesg,f)
    | Arg.Float f -> Float (mesg,f) | Arg.Set_float f -> Set_float (mesg,f)
    | Arg.Tuple f -> Tuple (mesg, List.map (spec_of_legacy_aux "") f) 
    | Arg.Symbol (l,f) -> Symbol (mesg,l,f)
    | Arg.Rest f -> Rest f
    | Arg.String f -> File (mesg,None,f)
    | Arg.Set_string f -> Set_string (mesg,f)

  let spec_of_legacy mesg d = 
    let mesg = match mesg with
      | None -> ""
      | Some s -> s in
    spec_of_legacy_aux mesg d

  open Completion

  let compl_speclist all = fun (c,e,d,h) ->
    let mk ?(excl=[]) x = 
      { compl = x;
        options = [c];
        excluded = excl@e;
        help = h
      } in
    let rec aux = function
    | Unit _ | Set _ | Clear _ -> mk None
    | Bool (s,_) -> mk (Some (CSymbol ["true";"false"],s))
    | Int (s,_) | Set_int (s,_) -> mk (Some (CInt,s))
    | Float (s,_) | Set_float (s,_) -> mk (Some (CFloat,s))
    | Tuple (s,l) -> let f x = 
        match (aux x).compl with
          | None -> invalid_arg "Rest, Set, Unit and Clear shouldn't be in tuple"
          | Some (c,s) -> (c,s) in
      mk (Some (CTuple (List.map f l),s))
    | Symbol (m,l,_) -> mk (Some (CSymbol l,m))
    | Rest _ -> mk ~excl:(List.map (fun (c,_,_,_) -> c) all) None
    | String (s,_) | Set_string (s,_) -> mk (Some (CString,s))
    | File (m,s,_) | Set_file (m,s,_) -> mk (Some (CFile s,m))
    | Dir (m,s,_) | Set_dir (m,s,_) -> mk (Some (CDir s,m))
    in aux d
    
  let trad_speclist = List.map 
      (fun (c,_,d,h) -> (c,legacy_of_spec d,h))


  let parse_argv_aux ?current argv speclist anonfun errmsg =
    parse_argv ?current argv (trad_speclist speclist) anonfun errmsg

  let parse_aux speclist anonfun errmsg =
      parse (trad_speclist speclist) anonfun errmsg

  let usage_aux speclist errmsg =
      usage (trad_speclist speclist) errmsg

  (** Completion Part *)
  let gen_zsh = ref false

  let speclist_compl = lazy 
    ["--gen_zsh", ["--gen_zsh"], Set gen_zsh, "Trigger the generation of the zsh completion function"]
  let errmsg_compl = "Generation of the completion function for this program : one output must be selected"
  let anonfun_compl _ = usage_aux (Lazy.force speclist_compl) errmsg_compl

  let main_gen speclist rest =
    if not !gen_zsh then
      anonfun_compl ();
    let prog = { name = (Filename.basename Sys.argv.(0));
                 infos = [];
                 rest = rest;
                 descs = List.map (compl_speclist speclist) speclist} in
     let o = open_out (sprintf "_%s" prog.name) in
     Zsh.write prog o;
     close_out o
    
        
     (** Last Part *)

  let parse_argv ?current argv speclist (rest,anonfun) errmsg =
    if is_mode_compl then
      (parse_argv_aux ?current argv (Lazy.force speclist_compl) anonfun_compl errmsg_compl;
      main_gen speclist rest)
    else
      parse_argv_aux ?current argv speclist anonfun errmsg
     

  let parse speclist (rest,anonfun) errmsg =
    if is_mode_compl then
      (parse_aux (Lazy.force speclist_compl) anonfun_compl errmsg_compl;
       main_gen speclist rest)
    else
      parse_aux speclist anonfun errmsg

  let usage speclist errmsg =
    if is_mode_compl then
      usage_aux (Lazy.force speclist_compl) errmsg_compl
    else
      usage_aux speclist errmsg
end

   
module Arg_legacy = Arg 

module Arg =
struct
  include Arg_legacy

  let realign = ref []
    
  let realign_ h = try let pre,suf = List.assoc h !realign in
  (Some pre),suf
  with Not_found -> (None,h)

  let trad_spec_list_leg (c,d,h) = 
    let mesg,help = realign_ h in
    (c,[],Arg_compl.spec_of_legacy mesg d,help)

  let parse_argv ?current argv speclist anonfun errmsg =
    if is_mode_compl then
      (Arg_compl.parse_argv_aux ?current argv (Lazy.force Arg_compl.speclist_compl) 
         Arg_compl.anonfun_compl Arg_compl.errmsg_compl;
       Arg_compl.main_gen (List.map trad_spec_list_leg speclist) (Completion.CFile(None), "?"))
    else
      parse_argv ?current argv speclist anonfun errmsg

  let parse speclist anonfun errmsg =
    if is_mode_compl then
      (Arg_compl.parse_aux (Lazy.force Arg_compl.speclist_compl) 
           Arg_compl.anonfun_compl Arg_compl.errmsg_compl;
       Arg_compl.main_gen (List.map trad_spec_list_leg speclist) (Completion.CFile(None), "?"))
    else
      parse speclist anonfun errmsg

  let rec second_word s =
    let len = String.length s in
    let rec loop n =
      if n >= len then len
      else if s.[n] = ' ' then loop (n+1)
      else n
    in
    try loop (String.index s ' ')
    with Not_found -> len

  let split ksd =
    match ksd with
      | (kwd, Symbol _, msg) ->
          ("",msg)
      | (kwd, spec, msg) ->
          let cutcol = second_word msg in
          let prefix = String.sub msg 0 cutcol in
          let suffix = String.sub msg cutcol (String.length msg - cutcol) in
          (prefix,suffix)

  (* from Arg *)
  let rec assoc3 x l =
    match l with
      | [] -> raise Not_found
      | (y1, y2, y3) :: t when y1 = x -> y2
      | _ :: t -> assoc3 x t
  ;;

  let help_action () = ();;

  let add_help_dumb speclist =
    let add1 =
      try ignore (assoc3 "-help" speclist); []
      with Not_found ->
            ["-help", Unit help_action, " Display this list of options"]
    and add2 =
      try ignore (assoc3 "--help" speclist); []
      with Not_found ->
        ["--help", Unit help_action, " Display this list of options"]
    in
    speclist @ (add1 @ add2)
  ;;


  let align speclist_bef =
    if is_mode_compl then
      begin
        (* in order to have the same number of element*)
        let speclist_bef = add_help_dumb speclist_bef in 
        let speclist_after = align speclist_bef in
        realign := List.fold_left2 (fun l bef (_,_,after) -> (after,split bef)::l) !realign speclist_bef speclist_after;
      end;
    align speclist_bef
      
end