csv.ml

   1 (* csv.ml - comma separated values parser
   2  *
   3  * $Id: csv.ml,v 1.6 2005-08-13 10:10:31 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             ) else (                    (* Bad single quote in field. *)
 135               field := c :: '"' :: !field;
 136               state := InQuotedField
 137             )
 138     ); (* end of match *)
 139     loop ()
 140   in
 141   try
 142     loop ()
 143   with
 144       End_of_file ->
 145         (* Any part left to write out? *)
 146         (match !state with
 147              StartField ->
 148                if !row <> [] then
 149                  ( empty_field (); end_of_row () )
 150            | InUnquotedField | InQuotedFieldAfterQuote ->
 151                end_of_field (); end_of_row ()
 152            | InQuotedField ->
 153                raise (Bad_CSV_file "Missing end quote after quoted field.")
 154         )
 155
 156 let load_in ?separator chan =
 157   let csv = ref [] in
 158   let f row =
 159     csv := row :: !csv
 160   in
 161   load_rows ?separator f chan;
 162   List.rev !csv
 163
 164 let load ?separator filename =
 165   let chan = open_in filename in
 166   let csv = load_in ?separator chan in
 167   close_in chan;
 168   csv
 169
 170 let trim ?(top=true) ?(left=true) ?(right=true) ?(bottom=true) csv =
 171   let rec empty_row = function
 172     | [] -> true
 173     | x :: xs when x <> "" -> false
 174     | x :: xs -> empty_row xs
 175   in
 176   let csv = if top then dropwhile empty_row csv else csv in
 177   let csv =
 178     if right then
 179       List.map (fun row ->
 180                   let row = List.rev row in
 181                   let row = dropwhile ((=) "") row in
 182                   let row = List.rev row in
 183                   row) csv
 184     else csv in
 185   let csv =
 186     if bottom then (
 187       let csv = List.rev csv in
 188       let csv = dropwhile empty_row csv in
 189       let csv = List.rev csv in
 190       csv
 191     ) else csv in
 192
 193   let empty_left_cell =
 194     function [] -> true | x :: xs when x = "" -> true | _ -> false in
 195   let empty_left_col =
 196     List.fold_left (fun a row -> a && empty_left_cell row) true in
 197   let remove_left_col =
 198     List.map (function [] -> [] | x :: xs -> xs) in
 199   let rec loop csv =
 200     if empty_left_col csv then (
 201       let csv = remove_left_col csv in
 202       loop csv
 203     ) else csv
 204   in
 205
 206   let csv = if left then loop csv else csv in
 207
 208   csv
 209
 210 let square csv =
 211   let columns = columns csv in
 212   List.map (
 213     fun row ->
 214       let n = List.length row in
 215       let row = List.rev row in
 216       let rec loop acc = function
 217         | 0 -> acc
 218         | i -> "" :: loop acc (i-1)
 219       in
 220       let row = loop row (columns - n) in
 221       List.rev row
 222   ) csv
 223
 224 let associate header data =
 225   let nr_cols = List.length header in
 226   let rec trunc = function
 227     | 0, _ -> []
 228     | n, [] -> "" :: trunc (n-1, [])
 229     | n, (x :: xs) -> x :: trunc (n-1, xs)
 230   in
 231   List.map (
 232     fun row ->
 233       let row = trunc (nr_cols, row) in
 234       List.combine header row
 235   ) data
 236
 237 let save_out ?(separator = ',') chan csv =
 238   (* Quote a single CSV field. *)
 239   let quote_field field =
 240     if String.contains field separator ||
 241       String.contains field '\"' ||
 242       String.contains field '\n'
 243     then (
 244       let buffer = Buffer.create 100 in
 245       Buffer.add_char buffer '\"';
 246       for i = 0 to (String.length field) - 1 do
 247         match field.[i] with
 248             '\"' -> Buffer.add_string buffer "\"\""
 249           | c    -> Buffer.add_char buffer c
 250       done;
 251       Buffer.add_char buffer '\"';
 252       Buffer.contents buffer
 253     )
 254     else
 255       field
 256   in
 257
 258   let separator = String.make 1 separator in
 259   List.iter (fun line ->
 260                output_string chan (String.concat separator
 261                                      (List.map quote_field line));
 262                output_char chan '\n') csv
 263
 264 let print ?separator csv =
 265   save_out ?separator stdout csv; flush stdout
 266
 267 let save ?separator file csv =
 268   let chan = open_out file in
 269   save_out ?separator chan csv;
 270   close_out chan
 271
 272 let save_out_readable chan csv =
 273   (* Escape all the strings in the CSV file first. *)
 274   let csv = List.map (List.map String.escaped) csv in
 275
 276   let csv = square csv in
 277
 278   (* Find the width of each column. *)
 279   let widths =
 280     match csv with
 281       | [] -> []
 282       | r :: _ ->
 283           let n = List.length r in
 284           let lengths = List.map (List.map String.length) csv in
 285           let max2rows r1 r2 =
 286             let rp = List.combine r1 r2 in
 287             List.map (fun ((a : int), (b : int)) -> max a b) rp
 288           in
 289           let rec repeat x = function
 290             | 0 -> []
 291             | i -> x :: repeat x (i-1)
 292           in
 293           List.fold_left max2rows (repeat 0 n) lengths in
 294
 295   (* Print out each cell at the correct width. *)
 296   let rec repeat f = function
 297     | 0 -> ()
 298     | i -> f (); repeat f (i-1)
 299   in
 300   List.iter (
 301     fun row ->
 302       let row = List.combine widths row in
 303       List.iter (
 304         fun (width, cell) ->
 305           output_string chan cell;
 306           let n = String.length cell in
 307           repeat (fun () -> output_char chan ' ') (width - n + 1)
 308       ) row;
 309       output_char chan '\n'
 310   ) csv
 311
 312 let print_readable = save_out_readable stdout