csv.ml

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