csvtool.ml

   1 (* Handy tool for managing CSV files.
   2  * $Id: csvtool.ml,v 1.5 2006-10-24 10:09:33 rich Exp $
   3  *)
   4
   5 open Printf
   6 open Csv
   7
   8 (*------------------------------ start of code from extlib *)
   9 exception Invalid_string
  10
  11 let find str sub =
  12   let sublen = String.length sub in
  13   if sublen = 0 then
  14     0
  15   else
  16     let found = ref 0 in
  17     let len = String.length str in
  18     try
  19       for i = 0 to len - sublen do
  20         let j = ref 0 in
  21         while String.unsafe_get str (i + !j) = String.unsafe_get sub !j do
  22           incr j;
  23           if !j = sublen then begin found := i; raise Exit; end;
  24         done;
  25       done;
  26       raise Invalid_string
  27     with
  28       Exit -> !found
  29
  30 let split str sep =
  31   let p = find str sep in
  32   let len = String.length sep in
  33   let slen = String.length str in
  34   String.sub str 0 p, String.sub str (p + len) (slen - p - len)
  35
  36 let nsplit str sep =
  37   if str = "" then []
  38   else (
  39     let rec nsplit str sep =
  40       try
  41         let s1 , s2 = split str sep in
  42         s1 :: nsplit s2 sep
  43       with
  44         Invalid_string -> [str]
  45     in
  46     nsplit str sep
  47   )
  48
  49 type 'a mut_list =  {
  50         hd: 'a;
  51         mutable tl: 'a list
  52 }
  53 external inj : 'a mut_list -> 'a list = "%identity"
  54
  55 let dummy_node () = { hd = Obj.magic (); tl = [] }
  56
  57 let rec drop n = function
  58   | _ :: l when n > 0 -> drop (n-1) l
  59   | l -> l
  60
  61 let take n l =
  62   let rec loop n dst = function
  63     | h :: t when n > 0 ->
  64         let r = { hd = h; tl = [] } in
  65         dst.tl <- inj r;
  66         loop (n-1) r t
  67     | _ ->
  68         ()
  69   in
  70   let dummy = dummy_node() in
  71   loop n dummy l;
  72   dummy.tl
  73 (*------------------------------ end of extlib code *)
  74
  75 (* Parse column specs. *)
  76 type colspec = range list
  77 and range =
  78   | Col of int (* 0 *)
  79   | Range of int * int (* 2-5 *)
  80   | ToEnd of int (* 7- *)
  81
  82 let parse_colspec ~count_zero colspec =
  83   let cols = nsplit colspec "," in
  84   let cols = List.map (
  85     fun col ->
  86       try
  87         (try
  88            let first, second = split col "-" in
  89            if second <> "" then
  90              Range (int_of_string first, int_of_string second)
  91            else
  92              ToEnd (int_of_string first)
  93          with
  94            Invalid_string ->
  95              Col (int_of_string col)
  96         )
  97       with
  98         Failure "int_of_string" ->
  99           failwith (colspec ^ ":" ^ col ^ ": invalid column-spec")
 100   ) cols in
 101
 102   (* Adjust so columns always count from zero. *)
 103   if not count_zero then
 104     List.map (
 105       function
 106       | Col c -> Col (c-1)
 107       | Range (s, e) -> Range (s-1, e-1)
 108       | ToEnd e -> ToEnd (e-1)
 109     ) cols
 110   else
 111     cols
 112
 113 let rec width_of_colspec = function
 114   | [] -> 0
 115   | Col c :: rest -> 1 + width_of_colspec rest
 116   | Range (s, e) :: rest -> (e-s+1) + width_of_colspec rest
 117   | ToEnd _ :: _ ->
 118       failwith "width_of_colspec: cannot calculate width of an open column spec (one which contains 'N-')"
 119
 120 (* For closed column specs, this preserves the correct width in the
 121  * result.
 122  *)
 123 let cols_of_colspec colspec row =
 124   let rec loop = function
 125     | [] -> []
 126     | Col c :: rest ->
 127         (try List.nth row c
 128          with Failure "nth" -> "") :: loop rest
 129     | Range (s, e) :: rest ->
 130         let width = e-s+1 in
 131         let range = take width (drop s row) in
 132         let range = List.hd (set_columns width [range]) in
 133         List.append range (loop rest)
 134     | ToEnd e :: rest ->
 135         List.append (drop e row) (loop rest)
 136   in
 137   loop colspec
 138
 139 (* The actual commands. *)
 140 let cmd_cols ~input_sep ~output_sep ~chan colspec files =
 141   List.iter (
 142     fun filename ->
 143       let csv = load ~separator:input_sep filename in
 144       let csv = List.map (cols_of_colspec colspec) csv in
 145       save_out ~separator:output_sep chan csv
 146   ) files
 147
 148 let cmd_namedcols ~input_sep ~output_sep ~chan names files =
 149   List.iter (
 150     fun filename ->
 151       let csv = load ~separator:input_sep filename in
 152       let header, data =
 153         match csv with
 154         | [] -> failwith "no rows in this CSV file"
 155         | h :: t -> h, t in
 156       let data = associate header data in
 157       let data = List.map (
 158         fun row -> List.map (fun name -> List.assoc name row) names
 159       ) data in
 160       save_out ~separator:output_sep chan data
 161   ) files
 162
 163 let cmd_width ~input_sep ~chan files =
 164   let width = List.fold_left (
 165     fun width filename ->
 166       let csv = load ~separator:input_sep filename in
 167       let width = max width (columns csv) in
 168       width
 169   ) 0 files in
 170   fprintf chan "%d\n" width
 171
 172 let cmd_height ~input_sep ~chan files =
 173   let height = List.fold_left (
 174     fun height filename ->
 175       let csv = load ~separator:input_sep filename in
 176       let height = height + lines csv in
 177       height
 178   ) 0 files in
 179   fprintf chan "%d\n" height
 180
 181 let cmd_readable ~input_sep ~chan files =
 182   let csv = List.concat (List.map (load ~separator:input_sep) files) in
 183   save_out_readable chan csv
 184
 185 let cmd_cat ~input_sep ~output_sep ~chan files =
 186   (* Avoid loading the whole file into memory. *)
 187   let f row =
 188     save_out ~separator:output_sep chan [row]
 189   in
 190   List.iter (
 191     fun filename ->
 192       let in_chan = open_in filename in
 193       load_rows ~separator:input_sep f in_chan;
 194       close_in in_chan
 195   ) files
 196
 197 let cmd_square ~input_sep ~output_sep ~chan files =
 198   let csv = List.concat (List.map (load ~separator:input_sep) files) in
 199   let csv = square csv in
 200   save_out ~separator:output_sep chan csv
 201
 202 let cmd_sub ~input_sep ~output_sep ~chan r c rows cols files =
 203   let csv = List.concat (List.map (load ~separator:input_sep) files) in
 204   let csv = sub r c rows cols csv in
 205   save_out ~separator:output_sep chan csv
 206
 207 let cmd_replace ~input_sep ~output_sep ~chan colspec update files =
 208   let csv = List.concat (List.map (load ~separator:input_sep) files) in
 209
 210   (* Load the update CSV file in. *)
 211   let update = Csv.load ~separator:input_sep update in
 212
 213   (* Compare two rows for equality by considering only the columns
 214    * in colspec.
 215    *)
 216   let equal row1 row2 =
 217     let row1 = cols_of_colspec colspec row1 in
 218     let row2 = cols_of_colspec colspec row2 in
 219     0 = Csv.compare [row1] [row2]
 220   in
 221
 222   (* Look for rows in the original to be replaced by rows from the
 223    * update file.  This is an ugly O(n^2) hack (XXX).
 224    *)
 225   let csv = List.filter (
 226     fun row -> not (List.exists (equal row) update)
 227   ) csv in
 228   let csv = csv @ update in
 229   save_out ~separator:output_sep chan csv
 230
 231 let rec uniq = function
 232   | [] -> []
 233   | [x] -> [x]
 234   | x :: y :: xs when Pervasives.compare x y = 0 ->
 235       uniq (x :: xs)
 236   | x :: y :: xs ->
 237       x :: uniq (y :: xs)
 238
 239 let cmd_join ~input_sep ~output_sep ~chan colspec1 colspec2 files =
 240   (* Load in the files separately. *)
 241   let csvs = List.map (load ~separator:input_sep) files in
 242
 243   (* For each CSV file, construct a hash table from row class (key) to
 244    * the (possibly empty) output columns (values).
 245    * Also construct a hash which has the unique list of row classes.
 246    *)
 247   let keys = Hashtbl.create 1023 in
 248   let hashes = List.map (
 249     fun csv ->
 250       let hash = Hashtbl.create 1023 in
 251       List.iter (
 252         fun row ->
 253           let key = cols_of_colspec colspec1 row in
 254           let value = cols_of_colspec colspec2 row in
 255           if not (Hashtbl.mem keys key) then Hashtbl.add keys key true;
 256           Hashtbl.add hash key value
 257       ) csv;
 258       hash
 259   ) csvs in
 260
 261   (* Get the keys. *)
 262   let keys = Hashtbl.fold (fun key _ xs -> key :: xs) keys [] in
 263
 264   let value_width = width_of_colspec colspec2 in
 265   let empty_value =
 266     List.hd (set_columns value_width [[""]]) in
 267   let multiple_values =
 268     List.hd (set_columns value_width [["!MULTIPLE VALUES"]]) in
 269
 270   (* Generate output CSV. *)
 271   let keys = List.sort Pervasives.compare keys in
 272   let keys = List.map (fun key -> key, []) keys in
 273   let csv = List.fold_left (
 274     fun keys hash ->
 275       List.map (
 276         fun (key, values) ->
 277           let value = try Hashtbl.find_all hash key with Not_found -> [] in
 278           let value =
 279             match value with
 280             | [] -> empty_value
 281             | [value] -> value
 282             | _::_ -> multiple_values in
 283           key, (value :: values)
 284       ) keys
 285   ) keys hashes in
 286   let csv = List.map (
 287     fun (key, values) ->
 288       key @ List.flatten (List.rev values)
 289   ) csv in
 290   save_out ~separator:output_sep chan csv
 291
 292 (* Process the arguments. *)
 293 let usage =
 294   "csvtool - Copyright (C) 2005-2006 Richard W.M. Jones, Merjis Ltd.
 295
 296 csvtool is a tool for performing manipulations on CSV files from shell scripts.
 297
 298 Summary:
 299   csvtool [-options] command [command-args] input.csv [input2.csv [...]]
 300
 301 Commands:
 302   col <column-spec>
 303     Return one or more columns from the CSV file.  Columns are numbered
 304     starting from zero.
 305
 306     For <column-spec>, see below.
 307
 308       Example: csvtool col 1-3,6 input.csv > output.csv
 309
 310   namedcol <names>
 311     Assuming the first row of the CSV file is a list of column headings,
 312     this returned the column(s) with the named headings.
 313
 314     <names> is a comma-separated list of names.
 315
 316       Example: csvtool namedcol Account,Cost input.csv > output.csv
 317
 318   width
 319     Print the maximum width of the CSV file (number of columns in the
 320     widest row).
 321
 322   height
 323     Print the number of rows in the CSV file.
 324
 325   readable
 326     Print the input CSV in a readable format.
 327
 328   cat
 329     This concatenates the input files together and writes them to
 330     the output.  You can use this to change the separator character.
 331
 332       Example: csvtool -t TAB -u , cat input.tsv > output.csv
 333
 334   join <column-spec1> <column-spec2>
 335     Join (collate) multiple CSV files together.
 336
 337     <column-spec1> controls which columns are compared.
 338
 339     <column-spec2> controls which columns are copied into the new file.
 340
 341       Example:
 342         csvtool join 1 2 coll1.csv coll2.csv
 343         If coll1.csv contains:
 344           Computers,$40
 345           Software,$100
 346         and coll2.csv contains:
 347           Computers,$50
 348         then the output will be:
 349           Computers,$40,$50
 350           Software,$100,
 351
 352   square
 353     Make the CSV square, so all rows have the same length.
 354
 355   sub r c rows cols
 356     Take a square subset of the CSV, top left at row r, column c (counting
 357     from 0), which is rows deep and cols wide.
 358
 359   replace <column-spec> update.csv original.csv
 360     Replace rows in original.csv with rows from update.csv.  The columns
 361     in <column-spec> only are used to compare rows in input.csv and
 362     update.csv to see if they are candidates for replacement.
 363
 364       Example:
 365         csvtool replace 3 updates.csv original.csv > new.csv
 366         mv new.csv original.csv
 367
 368 Column specs:
 369   A <column-spec> is a comma-separated list of column numbers
 370   or column ranges.
 371
 372     Examples:
 373       1                       Column 1 (the first, leftmost column)
 374       2,5,7                   Columns 2, 5 and 7
 375       1-3,5                   Columns 1, 2, 3 and 5
 376       1,5-                    Columns 1, 5 and up.
 377
 378   Columns are numbered starting from 1 unless the -z option is given.
 379
 380 Input files:
 381   csvtool takes a list of input file(s) from the command line.
 382
 383 Output file:
 384   Normally the output is written to stdout.  Use the -o option
 385   to override this.
 386
 387 Separators:
 388   The default separator character is , (comma).  To change this
 389   on input or output see the -t and -u options respectively.
 390
 391   Use -t TAB or -u TAB (literally T-A-B!) to specify tab-separated
 392   files.
 393
 394 Options:"
 395
 396 let () =
 397   let input_sep = ref ',' in
 398   let set_input_sep = function
 399     | "TAB" -> input_sep := '\t'
 400     | "COMMA" -> input_sep := ','
 401     | "SPACE" -> input_sep := ' '
 402     | s -> input_sep := s.[0]
 403   in
 404
 405   let output_sep = ref ',' in
 406   let set_output_sep = function
 407     | "TAB" -> output_sep := '\t'
 408     | "COMMA" -> output_sep := ','
 409     | "SPACE" -> output_sep := ' '
 410     | s -> output_sep := s.[0]
 411   in
 412
 413   let count_zero = ref false in
 414
 415   let output_file = ref "" in
 416
 417   let argspec = [
 418     "-t", Arg.String set_input_sep,
 419     "Input separator char.  Use -t TAB for tab separated input.";
 420     "-u", Arg.String set_output_sep,
 421     "Output separator char.  Use -u TAB for tab separated output.";
 422     "-o", Arg.Set_string output_file,
 423     "Write output to file (instead of stdout)";
 424     "-z", Arg.Set count_zero,
 425     "Number columns from 0 instead of 1";
 426   ] in
 427
 428   let rest = ref [] in
 429   let set_rest str =
 430     rest := str :: !rest
 431   in
 432
 433   Arg.parse argspec set_rest usage;
 434
 435   let input_sep = !input_sep in
 436   let output_sep = !output_sep in
 437   let count_zero = !count_zero in
 438   let output_file = !output_file in
 439   let rest = List.rev !rest in
 440
 441   (* Set up the output file. *)
 442   let chan =
 443     if output_file <> "" then open_out output_file
 444     else stdout in
 445
 446   (match rest with
 447      | ("col"|"cols") :: colspec :: files ->
 448          let colspec = parse_colspec ~count_zero colspec in
 449          cmd_cols ~input_sep ~output_sep ~chan colspec files
 450      | ("namedcol"|"namedcols") :: names :: files ->
 451          let names = nsplit names "," in
 452          cmd_namedcols ~input_sep ~output_sep ~chan names files
 453      | "width" :: files ->
 454          cmd_width ~input_sep ~chan files
 455      | "height" :: files ->
 456          cmd_height ~input_sep ~chan files
 457      | "readable" :: files ->
 458          cmd_readable ~input_sep ~chan files
 459      | ("cat"|"concat") :: files ->
 460          cmd_cat ~input_sep ~output_sep ~chan files
 461      | ("join"|"collate") :: colspec1 :: colspec2 :: ((_::_::_) as files) ->
 462          let colspec1 = parse_colspec ~count_zero colspec1 in
 463          let colspec2 = parse_colspec ~count_zero colspec2 in
 464          cmd_join ~input_sep ~output_sep ~chan colspec1 colspec2 files
 465      | "square" :: files ->
 466          cmd_square ~input_sep ~output_sep ~chan files
 467      | "sub" :: r :: c :: rows :: cols :: files ->
 468          let r = int_of_string r in
 469          let c = int_of_string c in
 470          let rows = int_of_string rows in
 471          let cols = int_of_string cols in
 472          cmd_sub ~input_sep ~output_sep ~chan r c rows cols files
 473      | "replace" :: colspec :: update :: files ->
 474          let colspec = parse_colspec ~count_zero colspec in
 475          cmd_replace ~input_sep ~output_sep ~chan colspec update files
 476      | _ ->
 477          prerr_endline (Sys.executable_name ^ " --help for usage");
 478          exit 2
 479   );
 480
 481   if output_file <> "" then close_out chan