Read and write tab-delimited files.
[ocaml-csv.git] / csv.ml
1 (* csv.ml - comma separated values parser
2  *
3  * $Id: csv.ml,v 1.3 2004-12-22 13:47:51 rich Exp $
4  *)
5
6 (* The format of CSV files:
7  * 
8  * Each field starts with either a double quote char or some other
9  * char. For the some other char case things are simple: just read up
10  * to the next comma (,) which marks the end of the field.
11  * 
12  * In the case where a field begins with a double quote char the
13  * parsing rules are different. Any double quotes are doubled ("") and
14  * we finish reading when we reach an undoubled quote. eg: "The
15  * following is a quote: "", and that's all" is the CSV equivalent of
16  * the following literal field: The following is a quote: ", and that's
17  * all
18  *
19  * "0 is the quoted form of ASCII NUL.
20  * 
21  * CSV fields can also contain literal carriage return characters, if
22  * they are quoted, eg: "This field
23  * is split over lines" represents a
24  * single field containing a \n.
25  * 
26  * Excel will only use the quoting format if a field contains a double
27  * quote or comma, although there's no reason why Excel couldn't always
28  * use the quoted format.
29  * 
30  * The practical upshot of this is that you can't split a line in a CSV
31  * file just by looking at the commas. You need to parse each field
32  * separately.
33  * 
34  * How we represent CSV files:
35  * 
36  * We load in the whole CSV file at once, and store it internally as a
37  * 'string list list' type (note that each line in the CSV file can,
38  * and often will, have different lengths). We then provide simple
39  * functions to read the CSV file line-by-line, copy it out, or copy a
40  * subset of it into a matrix.
41  *)
42
43 type t = string list list
44
45 exception Bad_CSV_file of string
46
47 let rec dropwhile f = function
48   | [] -> []
49   | x :: xs when f x -> dropwhile f xs
50   | xs -> xs
51
52 let lines = List.length
53
54 let columns csv =
55   List.fold_left max 0 (List.map List.length csv)
56
57 type state_t = StartField
58                | InUnquotedField
59                | InQuotedField
60                | InQuotedFieldAfterQuote
61
62 let load_rows ?(separator = ',') f chan =
63   let row = ref [] in                   (* Current row. *)
64   let field = ref [] in                 (* Current field. *)
65   let state = ref StartField in         (* Current state. *)
66   let end_of_field () =
67     let field_list = List.rev !field in
68     let field_len = List.length field_list in
69     let field_str = String.create field_len in
70     let rec loop i = function
71         [] -> ()
72       | x :: xs ->
73           field_str.[i] <- x;
74           loop (i+1) xs
75     in
76     loop 0 field_list;
77     row := field_str :: !row;
78     field := [];
79     state := StartField
80   in
81   let empty_field () =
82     row := "" :: !row;
83     field := [];
84     state := StartField
85   in
86   let end_of_row () =
87     let row_list = List.rev !row in
88     f row_list;
89     row := [];
90     state := StartField
91   in
92   let rec loop () =
93     let c = input_char chan in
94     if c != '\r' then (                 (* Always ignore \r characters. *)
95       match !state with
96           StartField ->                 (* Expecting quote or other char. *)
97             if c = '\"' then (
98               state := InQuotedField;
99               field := []
100             ) else if c = separator then (* Empty field. *)
101               empty_field ()
102             else if c = '\n' then (     (* Empty field, end of row. *)
103               empty_field ();
104               end_of_row ()
105             ) else (
106               state := InUnquotedField;
107               field := [c]
108             )
109         | InUnquotedField ->            (* Reading chars to end of field. *)
110             if c = separator then       (* End of field. *)
111               end_of_field ()
112             else if c = '\n' then (     (* End of field and end of row. *)
113               end_of_field ();
114               end_of_row ()
115             ) else
116               field := c :: !field
117         | InQuotedField ->              (* Reading chars to end of field. *)
118             if c = '\"' then
119               state := InQuotedFieldAfterQuote
120             else
121               field := c :: !field
122         | InQuotedFieldAfterQuote ->
123             if c = '\"' then (          (* Doubled quote. *)
124               field := c :: !field;
125               state := InQuotedField
126             ) else if c = '0' then (    (* Quote-0 is ASCII NUL. *)
127               field := '\000' :: !field;
128               state := InQuotedField
129             ) else if c = separator then (* End of field. *)
130               end_of_field ()
131             else if c = '\n' then (     (* End of field and end of row. *)
132               end_of_field ();
133               end_of_row ()
134             )
135     ); (* end of match *)
136     loop ()
137   in
138   try
139     loop ()
140   with
141       End_of_file ->
142         (* Any part left to write out? *)
143         (match !state with
144              StartField ->
145                if !row <> [] then
146                  ( empty_field (); end_of_row () )
147            | InUnquotedField | InQuotedFieldAfterQuote ->
148                end_of_field (); end_of_row ()
149            | InQuotedField ->
150                raise (Bad_CSV_file "Missing end quote after quoted field.")
151         )
152
153 let load_in ?separator chan =
154   let csv = ref [] in
155   let f row =
156     csv := row :: !csv
157   in
158   load_rows ?separator f chan;
159   List.rev !csv
160
161 let load ?separator filename =
162   let chan = open_in filename in
163   let csv = load_in ?separator chan in
164   close_in chan;
165   csv 
166
167 let trim ?(top=true) ?(left=true) ?(right=true) ?(bottom=true) csv =
168   let rec empty_row = function
169     | [] -> true
170     | x :: xs when x <> "" -> false
171     | x :: xs -> empty_row xs
172   in
173   let csv = if top then dropwhile empty_row csv else csv in
174   let csv =
175     if right then
176       List.map (fun row ->
177                   let row = List.rev row in
178                   let row = dropwhile ((=) "") row in
179                   let row = List.rev row in
180                   row) csv
181     else csv in
182   let csv =
183     if bottom then (
184       let csv = List.rev csv in
185       let csv = dropwhile empty_row csv in
186       let csv = List.rev csv in
187       csv
188     ) else csv in
189
190   let empty_left_cell =
191     function [] -> true | x :: xs when x = "" -> true | _ -> false in
192   let empty_left_col =
193     List.fold_left (fun a row -> a && empty_left_cell row) true in
194   let remove_left_col =
195     List.map (function [] -> [] | x :: xs -> xs) in
196   let rec loop csv =
197     if empty_left_col csv then (
198       let csv = remove_left_col csv in
199       loop csv
200     ) else csv
201   in
202
203   let csv = if left then loop csv else csv in
204
205   csv
206
207 let associate header data =
208   let nr_cols = List.length header in
209   let rec trunc = function
210     | 0, _ -> []
211     | n, [] -> "" :: trunc (n-1, [])
212     | n, (x :: xs) -> x :: trunc (n-1, xs)
213   in
214   List.map (
215     fun row ->
216       let row = trunc (nr_cols, row) in
217       List.combine header row
218   ) data
219
220 let save_out ?(separator = ',') chan csv =
221   (* Quote a single CSV field. *)
222   let quote_field field =
223     if String.contains field separator ||
224       String.contains field '\"' ||
225       String.contains field '\n'
226     then (
227       let buffer = Buffer.create 100 in
228       Buffer.add_char buffer '\"';
229       for i = 0 to (String.length field) - 1 do
230         match field.[i] with
231             '\"' -> Buffer.add_string buffer "\"\""
232           | c    -> Buffer.add_char buffer c
233       done;
234       Buffer.add_char buffer '\"';
235       Buffer.contents buffer
236     )
237     else
238       field
239   in
240
241   let separator = String.make 1 separator in
242   List.iter (fun line ->
243                output_string chan (String.concat separator
244                                      (List.map quote_field line));
245                output_char chan '\n') csv
246
247 let print ?separator csv =
248   save_out ?separator stdout csv
249
250 let save ?separator file csv =
251   let chan = open_out file in
252   save_out ?separator chan csv;
253   close_out chan