csv.ml

   1 (* csv.ml - comma separated values parser
   2  *
   3  * $Id: csv.ml,v 1.12 2006-10-18 14:56:12 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 open Printf
  44
  45 type t = string list list
  46
  47 exception Bad_CSV_file of string
  48
  49 let rec dropwhile f = function
  50   | [] -> []
  51   | x :: xs when f x -> dropwhile f xs
  52   | xs -> xs
  53
  54 (* from extlib: *)
  55 let rec drop n = function
  56   | _ :: l when n > 0 -> drop (n-1) l
  57   | l -> l
  58
  59 let rec take n = function
  60   | x :: xs when n > 0 -> x :: take (pred n) xs
  61   | _ -> []
  62
  63 let lines = List.length
  64
  65 let columns csv =
  66   List.fold_left max 0 (List.map List.length csv)
  67
  68 type state_t = StartField
  69                | InUnquotedField
  70                | InQuotedField
  71                | InQuotedFieldAfterQuote
  72
  73 let load_rows ?(separator = ',') f chan =
  74   let row = ref [] in                   (* Current row. *)
  75   let field = ref [] in                 (* Current field. *)
  76   let state = ref StartField in         (* Current state. *)
  77   let end_of_field () =
  78     let field_list = List.rev !field in
  79     let field_len = List.length field_list in
  80     let field_str = String.create field_len in
  81     let rec loop i = function
  82         [] -> ()
  83       | x :: xs ->
  84           field_str.[i] <- x;
  85           loop (i+1) xs
  86     in
  87     loop 0 field_list;
  88     row := field_str :: !row;
  89     field := [];
  90     state := StartField
  91   in
  92   let empty_field () =
  93     row := "" :: !row;
  94     field := [];
  95     state := StartField
  96   in
  97   let end_of_row () =
  98     let row_list = List.rev !row in
  99     f row_list;
 100     row := [];
 101     state := StartField
 102   in
 103   let rec loop () =
 104     let c = input_char chan in
 105     if c != '\r' then (                 (* Always ignore \r characters. *)
 106       match !state with
 107           StartField ->                 (* Expecting quote or other char. *)
 108             if c = '"' then (
 109               state := InQuotedField;
 110               field := []
 111             ) else if c = separator then (* Empty field. *)
 112               empty_field ()
 113             else if c = '\n' then (     (* Empty field, end of row. *)
 114               empty_field ();
 115               end_of_row ()
 116             ) else (
 117               state := InUnquotedField;
 118               field := [c]
 119             )
 120         | InUnquotedField ->            (* Reading chars to end of field. *)
 121             if c = separator then       (* End of field. *)
 122               end_of_field ()
 123             else if c = '\n' then (     (* End of field and end of row. *)
 124               end_of_field ();
 125               end_of_row ()
 126             ) else
 127               field := c :: !field
 128         | InQuotedField ->              (* Reading chars to end of field. *)
 129             if c = '"' then
 130               state := InQuotedFieldAfterQuote
 131             else
 132               field := c :: !field
 133         | InQuotedFieldAfterQuote ->
 134             if c = '"' then (           (* Doubled quote. *)
 135               field := c :: !field;
 136               state := InQuotedField
 137             ) else if c = '0' then (    (* Quote-0 is ASCII NUL. *)
 138               field := '\000' :: !field;
 139               state := InQuotedField
 140             ) else if c = separator then (* End of field. *)
 141               end_of_field ()
 142             else if c = '\n' then (     (* End of field and end of row. *)
 143               end_of_field ();
 144               end_of_row ()
 145             ) else (                    (* Bad single quote in field. *)
 146               field := c :: '"' :: !field;
 147               state := InQuotedField
 148             )
 149     ); (* end of match *)
 150     loop ()
 151   in
 152   try
 153     loop ()
 154   with
 155       End_of_file ->
 156         (* Any part left to write out? *)
 157         (match !state with
 158              StartField ->
 159                if !row <> [] then
 160                  ( empty_field (); end_of_row () )
 161            | InUnquotedField | InQuotedFieldAfterQuote ->
 162                end_of_field (); end_of_row ()
 163            | InQuotedField ->
 164                raise (Bad_CSV_file "Missing end quote after quoted field.")
 165         )
 166
 167 let load_in ?separator chan =
 168   let csv = ref [] in
 169   let f row =
 170     csv := row :: !csv
 171   in
 172   load_rows ?separator f chan;
 173   List.rev !csv
 174
 175 let load ?separator filename =
 176   let chan = open_in filename in
 177   let csv = load_in ?separator chan in
 178   close_in chan;
 179   csv
 180
 181 let trim ?(top=true) ?(left=true) ?(right=true) ?(bottom=true) csv =
 182   let rec empty_row = function
 183     | [] -> true
 184     | x :: xs when x <> "" -> false
 185     | x :: xs -> empty_row xs
 186   in
 187   let csv = if top then dropwhile empty_row csv else csv in
 188   let csv =
 189     if right then
 190       List.map (fun row ->
 191                   let row = List.rev row in
 192                   let row = dropwhile ((=) "") row in
 193                   let row = List.rev row in
 194                   row) csv
 195     else csv in
 196   let csv =
 197     if bottom then (
 198       let csv = List.rev csv in
 199       let csv = dropwhile empty_row csv in
 200       let csv = List.rev csv in
 201       csv
 202     ) else csv in
 203
 204   let empty_left_cell =
 205     function [] -> true | x :: xs when x = "" -> true | _ -> false in
 206   let empty_left_col =
 207     List.fold_left (fun a row -> a && empty_left_cell row) true in
 208   let remove_left_col =
 209     List.map (function [] -> [] | x :: xs -> xs) in
 210   let rec loop csv =
 211     if empty_left_col csv then (
 212       let csv = remove_left_col csv in
 213       loop csv
 214     ) else csv
 215   in
 216
 217   let csv = if left then loop csv else csv in
 218
 219   csv
 220
 221 let square csv =
 222   let columns = columns csv in
 223   List.map (
 224     fun row ->
 225       let n = List.length row in
 226       let row = List.rev row in
 227       let rec loop acc = function
 228         | 0 -> acc
 229         | i -> "" :: loop acc (i-1)
 230       in
 231       let row = loop row (columns - n) in
 232       List.rev row
 233   ) csv
 234
 235 let is_square csv =
 236   let columns = columns csv in
 237   List.for_all (fun row -> List.length row = columns) csv
 238
 239 let rec set_columns cols = function
 240   | [] -> []
 241   | r :: rs ->
 242       let rec loop i cells =
 243         if i < cols then (
 244           match cells with
 245           | [] -> "" :: loop (succ i) []
 246           | c :: cs -> c :: loop (succ i) cs
 247         )
 248         else []
 249       in
 250       loop 0 r :: set_columns cols rs
 251
 252 let rec set_rows rows csv =
 253   if rows > 0 then (
 254     match csv with
 255     | [] -> [] :: set_rows (pred rows) []
 256     | r :: rs -> r :: set_rows (pred rows) rs
 257   )
 258   else []
 259
 260 let set_size rows cols csv =
 261   set_columns cols (set_rows rows csv)
 262
 263 let sub r c rows cols csv =
 264   let csv = drop r csv in
 265   let csv = List.map (drop c) csv in
 266   let csv = set_rows rows csv in
 267   let csv = set_columns cols csv in
 268   csv
 269
 270 (* Compare two rows for semantic equality - ignoring any blank cells
 271  * at the end of each row.
 272  *)
 273 let rec compare_row (row1 : string list) row2 =
 274   match row1, row2 with
 275   | [], [] -> 0
 276   | x :: xs, y :: ys ->
 277       let c = compare x y in
 278       if c <> 0 then c else compare_row xs ys
 279   | "" :: xs , [] ->
 280       compare_row xs []
 281   | x :: xs, [] ->
 282       1
 283   | [], "" :: ys ->
 284       compare_row [] ys
 285   | [], y :: ys ->
 286       -1
 287
 288 (* Semantic equality for CSV files. *)
 289 let rec compare (csv1 : t) csv2 =
 290   match csv1, csv2 with
 291   | [], [] -> 0
 292   | x :: xs, y :: ys ->
 293       let c = compare_row x y in
 294       if c <> 0 then c else compare xs ys
 295   | x :: xs, [] ->
 296       let c = compare_row x [] in
 297       if c <> 0 then c else compare xs []
 298   | [], y :: ys ->
 299       let c = compare_row [] y in
 300       if c <> 0 then c else compare [] ys
 301
 302 (* Concatenate - arrange left to right. *)
 303 let rec concat = function
 304   | [] -> []
 305   | [csv] -> csv
 306   | left_csv :: csvs ->
 307       (* Concatenate the remaining CSV files. *)
 308       let right_csv = concat csvs in
 309
 310       (* Set the height of the left and right CSVs to the same. *)
 311       let nr_rows = max (lines left_csv) (lines right_csv) in
 312       let left_csv = set_rows nr_rows left_csv in
 313       let right_csv = set_rows nr_rows right_csv in
 314
 315       (* Square off the left CSV. *)
 316       let left_csv = square left_csv in
 317
 318       (* Prepend the right CSV rows with the left CSV rows. *)
 319       List.map (
 320         fun (left_row, right_row) -> List.append left_row right_row
 321       ) (List.combine left_csv right_csv)
 322
 323 let to_array csv =
 324   Array.of_list (List.map Array.of_list csv)
 325
 326 let of_array csv =
 327   List.map Array.to_list (Array.to_list csv)
 328
 329 let associate header data =
 330   let nr_cols = List.length header in
 331   let rec trunc = function
 332     | 0, _ -> []
 333     | n, [] -> "" :: trunc (n-1, [])
 334     | n, (x :: xs) -> x :: trunc (n-1, xs)
 335   in
 336   List.map (
 337     fun row ->
 338       let row = trunc (nr_cols, row) in
 339       List.combine header row
 340   ) data
 341
 342 let save_out ?(separator = ',') chan csv =
 343   (* Quote a single CSV field. *)
 344   let quote_field field =
 345     if String.contains field separator ||
 346       String.contains field '\"' ||
 347       String.contains field '\n'
 348     then (
 349       let buffer = Buffer.create 100 in
 350       Buffer.add_char buffer '\"';
 351       for i = 0 to (String.length field) - 1 do
 352         match field.[i] with
 353             '\"' -> Buffer.add_string buffer "\"\""
 354           | c    -> Buffer.add_char buffer c
 355       done;
 356       Buffer.add_char buffer '\"';
 357       Buffer.contents buffer
 358     )
 359     else
 360       field
 361   in
 362
 363   let separator = String.make 1 separator in
 364   List.iter (fun line ->
 365                output_string chan (String.concat separator
 366                                      (List.map quote_field line));
 367                output_char chan '\n') csv
 368
 369 let print ?separator csv =
 370   save_out ?separator stdout csv; flush stdout
 371
 372 let save ?separator file csv =
 373   let chan = open_out file in
 374   save_out ?separator chan csv;
 375   close_out chan
 376
 377 let save_out_readable chan csv =
 378   (* Escape all the strings in the CSV file first. *)
 379   (* XXX Why are we doing this?  I commented it out anyway.
 380   let csv = List.map (List.map String.escaped) csv in
 381   *)
 382
 383   (* Find the width of each column. *)
 384   let widths =
 385     (* Don't consider rows with only a single element - typically
 386      * long titles.
 387      *)
 388     let csv = List.filter (function [_] -> false | _ -> true) csv in
 389
 390     (* Square the CSV file - makes the next step simpler to implement. *)
 391     let csv = square csv in
 392
 393     match csv with
 394       | [] -> []
 395       | row1 :: rest ->
 396           let lengths_row1 = List.map String.length row1 in
 397           let lengths_rest = List.map (List.map String.length) rest in
 398           let max2rows r1 r2 =
 399             let rp =
 400               try List.combine r1 r2
 401               with
 402                 Invalid_argument "List.combine" ->
 403                   failwith (sprintf "Csv.save_out_readable: internal error: length r1 = %d, length r2 = %d" (List.length r1) (List.length r2)) in
 404             List.map (fun ((a : int), (b : int)) -> max a b) rp
 405           in
 406           List.fold_left max2rows lengths_row1 lengths_rest in
 407
 408   (* Print out each cell at the correct width. *)
 409   let rec repeat f = function
 410     | 0 -> ()
 411     | i -> f (); repeat f (i-1)
 412   in
 413   List.iter (
 414     function
 415     | [cell] ->                         (* Single column. *)
 416         output_string chan cell;
 417         output_char chan '\n'
 418     | row ->                            (* Other. *)
 419         (* Pair up each cell with its max width. *)
 420         let row =
 421           let rec loop = function
 422             | ([], _) -> []
 423             | (_, []) -> failwith "Csv.save_out_readable: internal error"
 424             | (cell :: cells, width :: widths) ->
 425                 (cell, width) :: loop (cells, widths)
 426           in
 427           loop (row, widths) in
 428         List.iter (
 429           fun (cell, width) ->
 430             output_string chan cell;
 431             let n = String.length cell in
 432             repeat (fun () -> output_char chan ' ') (width - n + 1)
 433         ) row;
 434         output_char chan '\n'
 435   ) csv
 436
 437 let print_readable = save_out_readable stdout