tree

tree.ml

(* Copyright (C) 2014 marklrh
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>. *)

open Core.Std

module type TREE=sig
  type 'a tree
  val f_name: abs:string -> string
  val invisible: name:string -> bool
  val build_tree: s:string -> level:int -> string tree
  val reorder: folder:'a tree list -> 'a tree list
  val render: elist:bool list -> string
  val printer: elist:bool list -> name:string -> dir:bool -> unit
  val print_tree: tree:string tree -> e:bool list -> unit
  val print_trees: subtrees:string tree list -> ending:bool list -> unit
  val give_me_a_tree: unit -> string tree
  val tree_ls : unit -> unit
end;;


module Tree:TREE = struct

  (* Type representing UNIX and GNU/Linux file system *)
  type 'a tree = Node of 'a * int * 'a tree list | Leaf of 'a * int;;

  (* Extract dir name or filename : abs:string -> string *)
  let f_name ~abs = 
    try 
      match (String.rsplit2_exn abs ~on:'/') with
      | (_, name) -> name
    with 
    | Not_found -> "."

  (* Check if a directory is invisible *)
  let invisible ~name = 
    let dot = String.get (f_name ~abs:name) 0 in
    if dot = '.' then true else false

  (* Build tree: string -> string -> string tree list *)
  let rec build_tree ~s ~level =
    let is_dir = Sys.is_directory_exn ~follow_symlinks:false s in
    if not is_dir then Leaf(s, level) else
      let content = Sys.ls_dir s in
      let trees_map subtree =
        build_tree ~s:(s ^/ subtree) ~level:(level + 1) in
      let subtrees = List.map ~f:trees_map content in 
      Node(s, level, subtrees)

  (* Put directory at end *)
  let reorder ~folder =
    let is_leaf tree =
      match tree with
      | Leaf(_, _) -> true
      | _ -> false in
    let file,dir = List.partition_tf ~f:is_leaf folder in
    file @ dir

  (* Render, which draws branches properly *)
  let render ~elist = 
    let rec render_helper el acc = 
      match el with
      | [] -> acc
      | h :: r -> if h then render_helper r "    "^acc else 
                    render_helper r "│   "^acc in
    match elist with
    | [] -> ""
    | x :: rest -> let arrow = if x then "└── " else "├── " in 
                   let sep = render_helper rest "" in 
                   sep ^ arrow

  (* Printer *)
  let printer ~elist ~name ~dir=
    let absname = f_name ~abs:name in
    let br = render ~elist:elist in
    print_string ("\027[30m "^br); 
    if dir then 
      print_endline ("\027[31m "^absname) 
    else
      print_endline ("\027[32m "^absname)

  (* Print in tree structure *)
  let rec print_tree ~tree ~e=
    match tree with
    | Leaf(name, _) -> printer ~elist:e ~name:name ~dir:false
    | Node(name,_,strees) -> 
       if invisible ~name:name then () else 
         (printer ~elist:e ~name:name ~dir:true;
          print_trees ~subtrees:(reorder ~folder:strees) ~ending:e)
  and print_trees ~subtrees ~ending=
    match subtrees with
    | [] -> ()
    | [st] -> print_tree ~tree:st ~e:(true :: ending)
    | st :: rest -> print_tree ~tree:st ~e:(false :: ending); 
                    print_trees ~subtrees:rest ~ending:ending

  (* Get tree structure of current directory *)
  let give_me_a_tree ()  = build_tree ~s:(Sys.getcwd()) ~level:0

  (* All combined *)
  let tree_ls() = print_tree ~tree:(give_me_a_tree()) ~e:[]

end;;

let () = Tree.tree_ls(); print_string "\027[0m "