csvtool: implement the "paste" command.
[csv/csv.git] / examples / csvtool.ml
1 (* Handy tool for managing CSV files.
2    @author Richard Jones <rjones@redhat.com>
3 *)
4
5 open Printf
6
7 (*------------------------------ start of code from extlib *)
8 exception Invalid_string
9
10 let find str sub =
11   let sublen = String.length sub in
12   if sublen = 0 then
13     0
14   else
15     let found = ref 0 in
16     let len = String.length str in
17     try
18       for i = 0 to len - sublen do
19         let j = ref 0 in
20         while String.unsafe_get str (i + !j) = String.unsafe_get sub !j do
21           incr j;
22           if !j = sublen then begin found := i; raise Exit; end;
23         done;
24       done;
25       raise Invalid_string
26     with
27       Exit -> !found
28
29 let split str sep =
30   let p = find str sep in
31   let len = String.length sep in
32   let slen = String.length str in
33   String.sub str 0 p, String.sub str (p + len) (slen - p - len)
34
35 let nsplit str sep =
36   if str = "" then []
37   else (
38     let rec nsplit str sep =
39       try
40         let s1 , s2 = split str sep in
41         s1 :: nsplit s2 sep
42       with
43         Invalid_string -> [str]
44     in
45     nsplit str sep
46   )
47
48 type 'a mut_list =  {
49         hd: 'a;
50         mutable tl: 'a list
51 }
52 external inj : 'a mut_list -> 'a list = "%identity"
53
54 let dummy_node () = { hd = Obj.magic (); tl = [] }
55
56 let rec drop n = function
57   | _ :: l when n > 0 -> drop (n-1) l
58   | l -> l
59
60 let take n l =
61   let rec loop n dst = function
62     | h :: t when n > 0 ->
63         let r = { hd = h; tl = [] } in
64         dst.tl <- inj r;
65         loop (n-1) r t
66     | _ ->
67         ()
68   in
69   let dummy = dummy_node() in
70   loop n dummy l;
71   dummy.tl
72 (*------------------------------ end of extlib code *)
73
74 (* Parse column specs. *)
75 type colspec = range list
76 and range =
77   | Col of int (* 0 *)
78   | Range of int * int (* 2-5 *)
79   | ToEnd of int (* 7- *)
80
81 let parse_colspec ~count_zero colspec =
82   let cols = nsplit colspec "," in
83   let cols = List.map (
84     fun col ->
85       try
86         (try
87            let first, second = split col "-" in
88            if second <> "" then
89              Range (int_of_string first, int_of_string second)
90            else
91              ToEnd (int_of_string first)
92          with
93            Invalid_string ->
94              Col (int_of_string col)
95         )
96       with
97         Failure "int_of_string" ->
98           failwith (colspec ^ ":" ^ col ^ ": invalid column-spec")
99   ) cols in
100
101   (* Adjust so columns always count from zero. *)
102   if not count_zero then
103     List.map (
104       function
105       | Col c -> Col (c-1)
106       | Range (s, e) -> Range (s-1, e-1)
107       | ToEnd e -> ToEnd (e-1)
108     ) cols
109   else
110     cols
111
112 let rec width_of_colspec = function
113   | [] -> 0
114   | Col c :: rest -> 1 + width_of_colspec rest
115   | Range (s, e) :: rest -> (e-s+1) + width_of_colspec rest
116   | ToEnd _ :: _ ->
117       failwith "width_of_colspec: cannot calculate width of an open column spec (one which contains 'N-')"
118
119 (* For closed column specs, this preserves the correct width in the
120  * result.
121  *)
122 let cols_of_colspec colspec row =
123   let rec loop = function
124     | [] -> []
125     | Col c :: rest ->
126         (try List.nth row c
127          with Failure "nth" -> "") :: loop rest
128     | Range (s, e) :: rest ->
129         let width = e-s+1 in
130         let range = take width (drop s row) in
131         let range = List.hd (Csv.set_columns width [range]) in
132         List.append range (loop rest)
133     | ToEnd e :: rest ->
134         List.append (drop e row) (loop rest)
135   in
136   loop colspec
137
138 (* The actual commands. *)
139 let cmd_cols ~input_sep ~output_sep ~chan colspec files =
140   List.iter (
141     fun filename ->
142       let csv = Csv.load ~separator:input_sep filename in
143       let csv = List.map (cols_of_colspec colspec) csv in
144       Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
145   ) files
146
147 let cmd_namedcols ~input_sep ~output_sep ~chan names files =
148   List.iter (
149     fun filename ->
150       let csv = Csv.load ~separator:input_sep filename in
151       let header, data =
152         match csv with
153         | [] -> failwith "no rows in this CSV file"
154         | h :: t -> h, t in
155       (* Do the headers requested exist in the CSV file?  If not,
156        * throw an error.
157        *)
158       List.iter (
159         fun name ->
160           if not (List.mem name header) then
161             failwith ("namedcol: requested header not in CSV file: " ^ name)
162       ) names;
163       let data = Csv.associate header data in
164       let data = List.map (
165         fun row -> List.map (fun name -> List.assoc name row) names
166       ) data in
167       let data = names :: data in
168       Csv.output_all (Csv.to_channel ~separator:output_sep chan) data
169   ) files
170
171 let cmd_width ~input_sep ~chan files =
172   let width = List.fold_left (
173     fun width filename ->
174       let csv = Csv.load ~separator:input_sep filename in
175       let width = max width (Csv.columns csv) in
176       width
177   ) 0 files in
178   fprintf chan "%d\n" width
179
180 let cmd_height ~input_sep ~chan files =
181   let height = List.fold_left (
182     fun height filename ->
183       let csv = Csv.load ~separator:input_sep filename in
184       let height = height + Csv.lines csv in
185       height
186   ) 0 files in
187   fprintf chan "%d\n" height
188
189 let cmd_readable ~input_sep ~chan files =
190   let csv = List.concat (List.map (Csv.load ~separator:input_sep) files) in
191   Csv.save_out_readable chan csv
192
193 let cmd_cat ~input_sep ~output_sep ~chan files =
194   (* Avoid loading the whole file into memory. *)
195   let chan = Csv.to_channel ~separator:output_sep chan in
196   let f row =
197     Csv.output_record chan row
198   in
199   List.iter (
200     fun filename ->
201       let in_chan, close =
202         match filename with
203         | "-" -> stdin, false
204         | filename -> open_in filename, true in
205       Csv.iter f (Csv.of_channel ~separator:input_sep in_chan);
206       if close then close_in in_chan
207   ) files
208
209 let cmd_paste ~input_sep ~output_sep ~chan files =
210   (* Return the 1st row, concatenation of all 1st rows; whether all
211      CSV files are empty; and the CSV files without their 1st row. *)
212   let rec add_columns = function
213     | [] -> ([], true, []) (* empty CSV file list *)
214     | [] :: csvs -> (* exhausted the first CSV file *)
215        let row, empty, csvs = add_columns csvs in
216        (row, empty, [] :: csvs)
217     | (r :: csv0) :: csvs ->
218        let row, empty, csvs = add_columns csvs in
219        (r @ row, false, csv0 :: csvs) in
220   let rec paste_rows csvs final_csv =
221     let row, empty, csvs = add_columns csvs in
222     if empty then List.rev final_csv
223     else paste_rows csvs (row :: final_csv)
224   in
225   let csvs = List.map (Csv.load ~separator:input_sep) files in
226   let csv = paste_rows csvs [] in
227   Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
228
229 let cmd_set_columns ~input_sep ~output_sep ~chan cols files =
230   (* Avoid loading the whole file into memory. *)
231   let f row =
232     let csv = [row] in
233     let csv = Csv.set_columns cols csv in
234     Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
235   in
236   List.iter (
237     fun filename ->
238       let in_chan, close =
239         match filename with
240         | "-" -> stdin, false
241         | filename -> open_in filename, true in
242       Csv.iter f (Csv.of_channel ~separator:input_sep in_chan);
243       if close then close_in in_chan
244   ) files
245
246 let cmd_set_rows ~input_sep ~output_sep ~chan rows files =
247   let csv = List.concat (List.map (Csv.load ~separator:input_sep) files) in
248   let csv = Csv.set_rows rows csv in
249   Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
250
251 let cmd_head ~input_sep ~output_sep ~chan rows files =
252   (* Avoid loading the whole file into memory, or even loading
253    * later files.
254    *)
255   let nr_rows = ref rows in
256   let chan = Csv.to_channel ~separator:output_sep chan in
257   let f row =
258     if !nr_rows > 0 then (
259       decr nr_rows;
260       Csv.output_record chan row
261     )
262   in
263   List.iter (
264     fun filename ->
265       if !nr_rows > 0 then (
266         let in_chan, close =
267           match filename with
268           | "-" -> stdin, false
269           | filename -> open_in filename, true in
270         Csv.iter f (Csv.of_channel ~separator:input_sep in_chan);
271         if close then close_in in_chan
272       )
273   ) files
274
275 let cmd_drop ~input_sep ~output_sep ~chan rows files =
276   (* Avoid loading the whole file into memory. *)
277   let nr_rows = ref rows in
278   let chan = Csv.to_channel ~separator:output_sep chan in
279   let f row =
280     if !nr_rows = 0 then
281       Csv.output_record chan row
282     else
283       decr nr_rows
284   in
285   List.iter (
286     fun filename ->
287       let in_chan, close =
288         match filename with
289         | "-" -> stdin, false
290         | filename -> open_in filename, true in
291       Csv.iter f (Csv.of_channel ~separator:input_sep in_chan);
292       if close then close_in in_chan
293   ) files
294
295 let cmd_square ~input_sep ~output_sep ~chan files =
296   let csv = List.concat (List.map (Csv.load ~separator:input_sep) files) in
297   let csv = Csv.square csv in
298   Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
299
300 let cmd_sub ~input_sep ~output_sep ~chan r c rows cols files =
301   let csv = List.concat (List.map (Csv.load ~separator:input_sep) files) in
302   let csv = Csv.sub r c rows cols csv in
303   Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
304
305 let cmd_replace ~input_sep ~output_sep ~chan colspec update files =
306   let csv = List.concat (List.map (Csv.load ~separator:input_sep) files) in
307
308   (* Load the update CSV file in. *)
309   let update = Csv.load ~separator:input_sep update in
310
311   (* Compare two rows for equality by considering only the columns
312    * in colspec.
313    *)
314   let equal row1 row2 =
315     let row1 = cols_of_colspec colspec row1 in
316     let row2 = cols_of_colspec colspec row2 in
317     0 = Csv.compare [row1] [row2]
318   in
319
320   (* Look for rows in the original to be replaced by rows from the
321    * update file.  This is an ugly O(n^2) hack (XXX).
322    *)
323   let csv = List.filter (
324     fun row -> not (List.exists (equal row) update)
325   ) csv in
326   let csv = csv @ update in
327   Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
328
329 let cmd_call ~input_sep ~output_sep ~chan command files =
330   (* Avoid loading the whole file into memory. *)
331   let f row =
332     let cmd =
333       command ^ " " ^ String.concat " " (List.map Filename.quote row) in
334     let code = Sys.command cmd in
335     if code <> 0 then (
336       eprintf "%s: terminated with exit code %d\n" command code;
337       exit code
338     )
339   in
340   List.iter (
341     fun filename ->
342       let in_chan, close =
343         match filename with
344         | "-" -> stdin, false
345         | filename -> open_in filename, true in
346       Csv.iter f (Csv.of_channel ~separator:input_sep in_chan);
347       if close then close_in in_chan
348   ) files
349
350 let rec uniq = function
351   | [] -> []
352   | [x] -> [x]
353   | x :: y :: xs when Pervasives.compare x y = 0 ->
354       uniq (x :: xs)
355   | x :: y :: xs ->
356       x :: uniq (y :: xs)
357
358 let cmd_join ~input_sep ~output_sep ~chan colspec1 colspec2 files =
359   (* Load in the files separately. *)
360   let csvs = List.map (Csv.load ~separator:input_sep) files in
361
362   (* For each CSV file, construct a hash table from row class (key) to
363    * the (possibly empty) output columns (values).
364    * Also construct a hash which has the unique list of row classes.
365    *)
366   let keys = Hashtbl.create 1023 in
367   let hashes = List.map (
368     fun csv ->
369       let hash = Hashtbl.create 1023 in
370       List.iter (
371         fun row ->
372           let key = cols_of_colspec colspec1 row in
373           let value = cols_of_colspec colspec2 row in
374           if not (Hashtbl.mem keys key) then Hashtbl.add keys key true;
375           Hashtbl.add hash key value
376       ) csv;
377       hash
378   ) csvs in
379
380   (* Get the keys. *)
381   let keys = Hashtbl.fold (fun key _ xs -> key :: xs) keys [] in
382
383   let value_width = width_of_colspec colspec2 in
384   let empty_value =
385     List.hd (Csv.set_columns value_width [[""]]) in
386   let multiple_values =
387     List.hd (Csv.set_columns value_width [["!MULTIPLE VALUES"]]) in
388
389   (* Generate output CSV. *)
390   let keys = List.sort Pervasives.compare keys in
391   let keys = List.map (fun key -> key, []) keys in
392   let csv = List.fold_left (
393     fun keys hash ->
394       List.map (
395         fun (key, values) ->
396           let value = try Hashtbl.find_all hash key with Not_found -> [] in
397           let value =
398             match value with
399             | [] -> empty_value
400             | [value] -> value
401             | _::_ -> multiple_values in
402           key, (value :: values)
403       ) keys
404   ) keys hashes in
405   let csv = List.map (
406     fun (key, values) ->
407       key @ List.flatten (List.rev values)
408   ) csv in
409   Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
410
411 let rec cmd_trim ~input_sep ~output_sep ~chan (top, left, right, bottom) files =
412   let csv = List.concat (List.map (Csv.load ~separator:input_sep) files) in
413   let csv = Csv.trim ~top ~left ~right ~bottom csv in
414   Csv.output_all (Csv.to_channel ~separator:output_sep chan) csv
415
416 and trim_flags flags =
417   let set c =
418     try ignore (String.index flags c); true with Not_found -> false
419   in
420   let top = set 't' in
421   let left = set 'l' in
422   let right = set 'r' in
423   let bottom = set 'b' in
424   (top, left, right, bottom)
425
426 (* Process the arguments. *)
427 let usage =
428   "csvtool - Copyright (C) 2005-2006 Richard W.M. Jones, Merjis Ltd.
429
430 csvtool is a tool for performing manipulations on CSV files from shell scripts.
431
432 Summary:
433   csvtool [-options] command [command-args] input.csv [input2.csv [...]]
434
435 Commands:
436   col <column-spec>
437     Return one or more columns from the CSV file.
438
439     For <column-spec>, see below.
440
441       Example: csvtool col 1-3,6 input.csv > output.csv
442
443   namedcol <names>
444     Assuming the first row of the CSV file is a list of column headings,
445     this returned the column(s) with the named headings.
446
447     <names> is a comma-separated list of names.
448
449       Example: csvtool namedcol Account,Cost input.csv > output.csv
450
451   width
452     Print the maximum width of the CSV file (number of columns in the
453     widest row).
454
455   height
456     Print the number of rows in the CSV file.
457
458     For most CSV files this is equivalent to 'wc -l', but note that
459     some CSV files can contain a row which breaks over two (or more)
460     lines.
461
462   setcolumns cols
463     Set the number of columns to cols (this also makes the CSV file
464     square).  Any short rows are padding with blank cells.  Any
465     long rows are truncated.
466
467   setrows rows
468     'setrows n' sets the number of rows to 'n'.  If there are fewer
469     than 'n' rows in the CSV files, then empty blank lines are added.
470
471   head rows
472   take rows
473     'head n' and 'take n' (which are synonyms) take the first 'n'
474     rows.  If there are fewer than 'n' rows, padding is not added.
475
476   drop rows
477     Drop the first 'rows' rows and return the rest (if any).
478
479       Example:
480         To remove the headings from a CSV file with headings:
481           csvtool drop 1 input.csv > output.csv
482
483         To extract rows 11 through 20 from a file:
484           csvtool drop 10 input.csv | csvtool take 10 - > output.csv
485
486   cat
487     This concatenates the input files together and writes them to
488     the output.  You can use this to change the separator character.
489
490       Example: csvtool -t TAB -u COMMA cat input.tsv > output.csv
491
492   paste
493     Concatenate the columns of the files together and write them to the
494     output.
495
496       Example: csvtool paste input1.csv input2.csv > output.csv
497
498   join <column-spec1> <column-spec2>
499     Join (collate) multiple CSV files together.
500
501     <column-spec1> controls which columns are compared.
502
503     <column-spec2> controls which columns are copied into the new file.
504
505       Example:
506         csvtool join 1 2 coll1.csv coll2.csv > output.csv
507
508         In the above example, if coll1.csv contains:
509           Computers,$40
510           Software,$100
511         and coll2.csv contains:
512           Computers,$50
513         then the output will be:
514           Computers,$40,$50
515           Software,$100,
516
517   square
518     Make the CSV square, so all rows have the same length.
519
520       Example: csvtool square input.csv > input-square.csv
521
522   trim [tlrb]+
523     Trim empty cells at the top/left/right/bottom of the CSV file.
524
525       Example:
526         csvtool trim t input.csv    # trims empty rows at the top only
527         csvtool trim tb input.csv   # trims empty rows at the top & bottom
528         csvtool trim lr input.csv   # trims empty columns at left & right
529         csvtool trim tlrb input.csv # trims empty rows/columns all around
530
531   sub r c rows cols
532     Take a square subset of the CSV, top left at row r, column c, which
533     is rows deep and cols wide.  'r' and 'c' count from 1, or
534     from 0 if -z option is given.
535
536   replace <column-spec> update.csv original.csv
537     Replace rows in original.csv with rows from update.csv.  The columns
538     in <column-spec> only are used to compare rows in input.csv and
539     update.csv to see if they are candidates for replacement.
540
541       Example:
542         csvtool replace 3 updates.csv original.csv > new.csv
543         mv new.csv original.csv
544
545   call command
546     This calls the external command (or shell function) 'command'
547     followed by a parameter for each column in the CSV file.  The
548     external command is called once for each row in the CSV file.
549     If any command returns a non-zero exit code then the whole
550     program terminates.
551
552       Tip:
553         Use the shell command 'export -f funcname' to export
554         a shell function for use as a command.  Within the
555         function, use the positional parameters $1, $2, ...
556         to refer to the columns.
557
558       Example (with a shell function):
559         function test {
560           echo Column 1: $1
561           echo Column 2: $2
562         }
563         export -f test
564         csvtool call test my.csv
565
566         In the above example, if my.csv contains:
567           how,now
568           brown,cow
569         then the output is:
570           Column 1: how
571           Column 2: now
572           Column 1: brown
573           Column 2: cow
574
575   readable
576     Print the input CSV in a readable format.
577
578 Column specs:
579   A <column-spec> is a comma-separated list of column numbers
580   or column ranges.
581
582     Examples:
583       1                       Column 1 (the first, leftmost column)
584       2,5,7                   Columns 2, 5 and 7
585       1-3,5                   Columns 1, 2, 3 and 5
586       1,5-                    Columns 1, 5 and up.
587
588   Columns are numbered starting from 1 unless the -z option is given.
589
590 Input files:
591   csvtool takes a list of input file(s) from the command line.
592
593   If an input filename is '-' then take input from stdin.
594
595 Output file:
596   Normally the output is written to stdout.  Use the -o option
597   to override this.
598
599 Separators:
600   The default separator character is , (comma).  To change this
601   on input or output see the -t and -u options respectively.
602
603   Use -t TAB or -u TAB (literally T-A-B!) to specify tab-separated
604   files.
605
606 Options:"
607
608 let () =
609   let input_sep = ref ',' in
610   let set_input_sep = function
611     | "TAB" -> input_sep := '\t'
612     | "COMMA" -> input_sep := ','
613     | "SPACE" -> input_sep := ' '
614     | s -> input_sep := s.[0]
615   in
616
617   let output_sep = ref ',' in
618   let set_output_sep = function
619     | "TAB" -> output_sep := '\t'
620     | "COMMA" -> output_sep := ','
621     | "SPACE" -> output_sep := ' '
622     | s -> output_sep := s.[0]
623   in
624
625   let count_zero = ref false in
626
627   let output_file = ref "" in
628
629   let rest = ref [] in
630   let set_rest str =
631     rest := str :: !rest
632   in
633
634   let argspec = [
635     "-t", Arg.String set_input_sep,
636     "Input separator char.  Use -t TAB for tab separated input.";
637     "-u", Arg.String set_output_sep,
638     "Output separator char.  Use -u TAB for tab separated output.";
639     "-o", Arg.Set_string output_file,
640     "Write output to file (instead of stdout)";
641     "-z", Arg.Set count_zero,
642     "Number columns from 0 instead of 1";
643     "-", Arg.Unit (fun () -> set_rest "-"),
644     "" (* Hack to allow '-' for input from stdin. *)
645   ] in
646
647   Arg.parse argspec set_rest usage;
648
649   let input_sep = !input_sep in
650   let output_sep = !output_sep in
651   let count_zero = !count_zero in
652   let output_file = !output_file in
653   let rest = List.rev !rest in
654
655   (* Set up the output file. *)
656   let chan =
657     if output_file <> "" then open_out output_file
658     else stdout in
659
660   (match rest with
661      | ("col"|"cols") :: colspec :: files ->
662          let colspec = parse_colspec ~count_zero colspec in
663          cmd_cols ~input_sep ~output_sep ~chan colspec files
664      | ("namedcol"|"namedcols") :: names :: files ->
665          let names = nsplit names "," in
666          cmd_namedcols ~input_sep ~output_sep ~chan names files
667      | ("width"|"columns") :: files ->
668          cmd_width ~input_sep ~chan files
669      | ("height"|"rows") :: files ->
670          cmd_height ~input_sep ~chan files
671      | "readable" :: files ->
672          cmd_readable ~input_sep ~chan files
673      | ("cat"|"concat") :: files ->
674          cmd_cat ~input_sep ~output_sep ~chan files
675      | "paste" :: files ->
676          cmd_paste ~input_sep ~output_sep ~chan files
677      | ("join"|"collate") :: colspec1 :: colspec2 :: ((_::_::_) as files) ->
678          let colspec1 = parse_colspec ~count_zero colspec1 in
679          let colspec2 = parse_colspec ~count_zero colspec2 in
680          cmd_join ~input_sep ~output_sep ~chan colspec1 colspec2 files
681      | "square" :: files ->
682          cmd_square ~input_sep ~output_sep ~chan files
683      | "sub" :: r :: c :: rows :: cols :: files ->
684          let r = int_of_string r in
685          let r = if not count_zero then r-1 else r in
686          let c = int_of_string c in
687          let c = if not count_zero then c-1 else c in
688          let rows = int_of_string rows in
689          let cols = int_of_string cols in
690          cmd_sub ~input_sep ~output_sep ~chan r c rows cols files
691      | "replace" :: colspec :: update :: files ->
692          let colspec = parse_colspec ~count_zero colspec in
693          cmd_replace ~input_sep ~output_sep ~chan colspec update files
694      | ("setcolumns"|"set_columns"|"set-columns"|
695             "setcols"|"set_cols"|"set-cols") :: cols :: files ->
696          let cols = int_of_string cols in
697          cmd_set_columns ~input_sep ~output_sep ~chan cols files
698      | ("setrows"|"set_rows"|"set-rows") :: rows :: files ->
699          let rows = int_of_string rows in
700          cmd_set_rows ~input_sep ~output_sep ~chan rows files
701      | ("head"|"take") :: rows :: files ->
702          let rows = int_of_string rows in
703          cmd_head ~input_sep ~output_sep ~chan rows files
704      | "drop" :: rows :: files ->
705          let rows = int_of_string rows in
706          cmd_drop ~input_sep ~output_sep ~chan rows files
707      | "call" :: command :: files ->
708          cmd_call ~input_sep ~output_sep ~chan command files
709      | "trim" :: flags :: files ->
710          let flags = trim_flags flags in
711          cmd_trim ~input_sep ~output_sep ~chan flags files
712      | _ ->
713          prerr_endline (Sys.executable_name ^ " --help for usage");
714          exit 2
715   );
716
717   if output_file <> "" then close_out chan