csvtool.ml

   1 (* Handy tool for managing CSV files.
   2  * $Id: csvtool.ml,v 1.2 2005-11-25 14:06:58 rich Exp $
   3  *)
   4
   5 open Printf
   6 open Csv
   7
   8 let cmd_cols ~separator ~csv ~chan cols =
   9   let cols = List.map int_of_string cols in
  10
  11   let output = List.map (
  12     fun row ->
  13       let n = List.length row in
  14       let row = List.map (
  15         fun col_wanted ->
  16           if 0 <= col_wanted && col_wanted < n then
  17             List.nth row col_wanted
  18           else
  19             ""
  20       ) cols in
  21       row
  22   ) csv in
  23   save_out ~separator chan output
  24
  25 let cmd_namedcols ~separator ~csv ~chan names =
  26   let header, data =
  27     match csv with
  28       | [] -> failwith "no rows in this CSV file"
  29       | h :: t -> h, t in
  30   let data = associate header data in
  31   let data = List.map (
  32     fun row -> List.map (fun name -> List.assoc name row) names
  33   ) data in
  34   save_out ~separator chan data
  35
  36 let cmd_width ~csv ~chan () =
  37   fprintf chan "%d\n" (columns csv)
  38
  39 let cmd_height ~csv ~chan () =
  40   fprintf chan "%d\n" (lines csv)
  41
  42 let cmd_readable ~csv ~chan () =
  43   save_out_readable chan csv
  44
  45 let cmd_square ~separator ~csv ~chan () =
  46   let csv = square csv in
  47   save_out ~separator chan csv
  48
  49 let cmd_sub ~separator ~csv ~chan args =
  50   let r, c, rows, cols =
  51     match args with
  52     | [ r; c; rows; cols ] ->
  53         int_of_string r, int_of_string c,
  54         int_of_string rows, int_of_string cols
  55     | _ ->
  56         failwith "unknown arguments to 'sub' command" in
  57   let csv = sub r c rows cols csv in
  58   save_out ~separator chan csv
  59
  60 (* Process the arguments. *)
  61 let usage =
  62   "csvtool - Copyright (C) 2005 Richard W.M. Jones, Merjis Ltd.
  63
  64 csvtool is a tool for performing manipulations on CSV files from shell scripts.
  65
  66 Summary:
  67   csvtool [-options] command [command-args] < input.csv
  68
  69 Commands:
  70   col [col1] [col2] ...
  71     Return one or more columns from the CSV file.  Columns are numbered
  72     starting from zero.
  73
  74   namedcol [name1] [name2] ...
  75     Assuming the first row of the CSV file is a list of column headings,
  76     this returned the column(s) with the named headings.
  77
  78   width
  79     Return the maximum width of the CSV file (number of columns in the
  80     widest row).
  81
  82   height
  83     Return the number of rows in the CSV file.
  84
  85   readable
  86     Print the input CSV in a readable format.
  87
  88   square
  89     Make the CSV square, so all rows have the same length.
  90
  91   sub r c rows cols
  92     Take a square subset of the CSV, top left at row r, column c (counting
  93     from 0), which is rows deep and cols wide.
  94
  95 Input and output files:
  96   csvtool normally processes its input from stdin and writes its output
  97   to stdout.  Use the -i and -o options to override this behaviour.
  98
  99 Options:"
 100
 101 let () =
 102   let input_sep = ref ',' in
 103   let set_input_sep = function
 104     | "TAB" -> input_sep := '\t'
 105     | "COMMA" -> input_sep := ','
 106     | s -> input_sep := s.[0]
 107   in
 108
 109   let output_sep = ref ',' in
 110   let set_output_sep = function
 111     | "TAB" -> output_sep := '\t'
 112     | "COMMA" -> output_sep := ','
 113     | s -> output_sep := s.[0]
 114   in
 115
 116   let input_file = ref "" in
 117   let output_file = ref "" in
 118
 119   let argspec = [
 120     "-t", Arg.String set_input_sep,
 121     "Input separator char.  Use -t TAB for tab separated input.";
 122     "-u", Arg.String set_output_sep,
 123     "Output separator char.  Use -t TAB for tab separated output.";
 124     "-i", Arg.Set_string input_file,
 125     "Read CSV input from file (instead of stdin)";
 126     "-o", Arg.Set_string output_file,
 127     "Write output to file (instead of stdout)"
 128   ] in
 129
 130   let rest = ref [] in
 131   let set_rest str =
 132     rest := str :: !rest
 133   in
 134
 135   Arg.parse argspec set_rest usage;
 136
 137   let input_sep = !input_sep in
 138   let output_sep = !output_sep in
 139   let input_file = !input_file in
 140   let output_file = !output_file in
 141   let rest = List.rev !rest in
 142
 143   let cmd, args =
 144     match rest with
 145       | [] -> prerr_endline (Sys.executable_name ^ " --help for usage"); exit 1
 146       | h :: t -> h, t in
 147
 148   (* Read the input file. *)
 149   let input =
 150     if input_file <> "" then load ~separator:input_sep input_file
 151     else load_in ~separator:input_sep stdin in
 152
 153   (* Set up the output file. *)
 154   let chan =
 155     if output_file <> "" then open_out output_file
 156     else stdout in
 157
 158   (match cmd with
 159      | "col" | "cols" ->
 160          cmd_cols ~separator:output_sep ~csv:input ~chan args
 161      | "namedcol" | "namedcols" ->
 162          cmd_namedcols ~separator:output_sep ~csv:input ~chan args
 163      | "width" ->
 164          cmd_width ~csv:input ~chan ()
 165      | "height" ->
 166          cmd_height ~csv:input ~chan ()
 167      | "readable" ->
 168          cmd_readable ~csv:input ~chan ()
 169      | "square" ->
 170          cmd_square ~separator:output_sep ~csv:input ~chan ()
 171      | "sub" ->
 172          cmd_sub ~separator:output_sep ~csv:input ~chan args
 173      | _ -> prerr_endline (Sys.executable_name ^ " --help for usage")
 174   );
 175
 176   if output_file <> "" then close_out chan