csv.ml

   1 (* csv.ml - comma separated values parser
   2  *
   3  * $Id: csv.ml,v 1.7 2005-11-25 14:06:58 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 (* from extlib: *)
  53 let rec drop n = function
  54   | _ :: l when n > 0 -> drop (n-1) l
  55   | l -> l
  56
  57 let rec take n = function
  58   | x :: xs when n > 0 -> x :: take (pred n) xs
  59   | _ -> []
  60
  61 let lines = List.length
  62
  63 let columns csv =
  64   List.fold_left max 0 (List.map List.length csv)
  65
  66 type state_t = StartField
  67                | InUnquotedField
  68                | InQuotedField
  69                | InQuotedFieldAfterQuote
  70
  71 let load_rows ?(separator = ',') f chan =
  72   let row = ref [] in                   (* Current row. *)
  73   let field = ref [] in                 (* Current field. *)
  74   let state = ref StartField in         (* Current state. *)
  75   let end_of_field () =
  76     let field_list = List.rev !field in
  77     let field_len = List.length field_list in
  78     let field_str = String.create field_len in
  79     let rec loop i = function
  80         [] -> ()
  81       | x :: xs ->
  82           field_str.[i] <- x;
  83           loop (i+1) xs
  84     in
  85     loop 0 field_list;
  86     row := field_str :: !row;
  87     field := [];
  88     state := StartField
  89   in
  90   let empty_field () =
  91     row := "" :: !row;
  92     field := [];
  93     state := StartField
  94   in
  95   let end_of_row () =
  96     let row_list = List.rev !row in
  97     f row_list;
  98     row := [];
  99     state := StartField
 100   in
 101   let rec loop () =
 102     let c = input_char chan in
 103     if c != '\r' then (                 (* Always ignore \r characters. *)
 104       match !state with
 105           StartField ->                 (* Expecting quote or other char. *)
 106             if c = '"' then (
 107               state := InQuotedField;
 108               field := []
 109             ) else if c = separator then (* Empty field. *)
 110               empty_field ()
 111             else if c = '\n' then (     (* Empty field, end of row. *)
 112               empty_field ();
 113               end_of_row ()
 114             ) else (
 115               state := InUnquotedField;
 116               field := [c]
 117             )
 118         | InUnquotedField ->            (* Reading chars to end of field. *)
 119             if c = separator then       (* End of field. *)
 120               end_of_field ()
 121             else if c = '\n' then (     (* End of field and end of row. *)
 122               end_of_field ();
 123               end_of_row ()
 124             ) else
 125               field := c :: !field
 126         | InQuotedField ->              (* Reading chars to end of field. *)
 127             if c = '"' then
 128               state := InQuotedFieldAfterQuote
 129             else
 130               field := c :: !field
 131         | InQuotedFieldAfterQuote ->
 132             if c = '"' then (           (* Doubled quote. *)
 133               field := c :: !field;
 134               state := InQuotedField
 135             ) else if c = '0' then (    (* Quote-0 is ASCII NUL. *)
 136               field := '\000' :: !field;
 137               state := InQuotedField
 138             ) else if c = separator then (* End of field. *)
 139               end_of_field ()
 140             else if c = '\n' then (     (* End of field and end of row. *)
 141               end_of_field ();
 142               end_of_row ()
 143             ) else (                    (* Bad single quote in field. *)
 144               field := c :: '"' :: !field;
 145               state := InQuotedField
 146             )
 147     ); (* end of match *)
 148     loop ()
 149   in
 150   try
 151     loop ()
 152   with
 153       End_of_file ->
 154         (* Any part left to write out? *)
 155         (match !state with
 156              StartField ->
 157                if !row <> [] then
 158                  ( empty_field (); end_of_row () )
 159            | InUnquotedField | InQuotedFieldAfterQuote ->
 160                end_of_field (); end_of_row ()
 161            | InQuotedField ->
 162                raise (Bad_CSV_file "Missing end quote after quoted field.")
 163         )
 164
 165 let load_in ?separator chan =
 166   let csv = ref [] in
 167   let f row =
 168     csv := row :: !csv
 169   in
 170   load_rows ?separator f chan;
 171   List.rev !csv
 172
 173 let load ?separator filename =
 174   let chan = open_in filename in
 175   let csv = load_in ?separator chan in
 176   close_in chan;
 177   csv
 178
 179 let trim ?(top=true) ?(left=true) ?(right=true) ?(bottom=true) csv =
 180   let rec empty_row = function
 181     | [] -> true
 182     | x :: xs when x <> "" -> false
 183     | x :: xs -> empty_row xs
 184   in
 185   let csv = if top then dropwhile empty_row csv else csv in
 186   let csv =
 187     if right then
 188       List.map (fun row ->
 189                   let row = List.rev row in
 190                   let row = dropwhile ((=) "") row in
 191                   let row = List.rev row in
 192                   row) csv
 193     else csv in
 194   let csv =
 195     if bottom then (
 196       let csv = List.rev csv in
 197       let csv = dropwhile empty_row csv in
 198       let csv = List.rev csv in
 199       csv
 200     ) else csv in
 201
 202   let empty_left_cell =
 203     function [] -> true | x :: xs when x = "" -> true | _ -> false in
 204   let empty_left_col =
 205     List.fold_left (fun a row -> a && empty_left_cell row) true in
 206   let remove_left_col =
 207     List.map (function [] -> [] | x :: xs -> xs) in
 208   let rec loop csv =
 209     if empty_left_col csv then (
 210       let csv = remove_left_col csv in
 211       loop csv
 212     ) else csv
 213   in
 214
 215   let csv = if left then loop csv else csv in
 216
 217   csv
 218
 219 let square csv =
 220   let columns = columns csv in
 221   List.map (
 222     fun row ->
 223       let n = List.length row in
 224       let row = List.rev row in
 225       let rec loop acc = function
 226         | 0 -> acc
 227         | i -> "" :: loop acc (i-1)
 228       in
 229       let row = loop row (columns - n) in
 230       List.rev row
 231   ) csv
 232
 233 let is_square csv =
 234   let columns = columns csv in
 235   List.for_all (fun row -> List.length row = columns) csv
 236
 237 let rec set_columns cols = function
 238   | [] -> []
 239   | r :: rs ->
 240       let rec loop i cells =
 241         if i < cols then (
 242           match cells with
 243           | [] -> "" :: loop (succ i) []
 244           | c :: cs -> c :: loop (succ i) cs
 245         )
 246         else []
 247       in
 248       loop 0 r :: set_columns cols rs
 249
 250 let rec set_rows rows csv =
 251   if rows > 0 then (
 252     match csv with
 253     | [] -> [] :: set_rows (pred rows) []
 254     | r :: rs -> r :: set_rows (pred rows) rs
 255   )
 256   else []
 257
 258 let set_size rows cols csv =
 259   set_columns cols (set_rows rows csv)
 260
 261 let sub r c rows cols csv =
 262   let csv = drop r csv in
 263   let csv = List.map (drop c) csv in
 264   let csv = set_rows rows csv in
 265   let csv = set_columns cols csv in
 266   csv
 267
 268 let to_array csv =
 269   Array.of_list (List.map Array.of_list csv)
 270
 271 let of_array csv =
 272   List.map Array.to_list (Array.to_list csv)
 273
 274 let associate header data =
 275   let nr_cols = List.length header in
 276   let rec trunc = function
 277     | 0, _ -> []
 278     | n, [] -> "" :: trunc (n-1, [])
 279     | n, (x :: xs) -> x :: trunc (n-1, xs)
 280   in
 281   List.map (
 282     fun row ->
 283       let row = trunc (nr_cols, row) in
 284       List.combine header row
 285   ) data
 286
 287 let save_out ?(separator = ',') chan csv =
 288   (* Quote a single CSV field. *)
 289   let quote_field field =
 290     if String.contains field separator ||
 291       String.contains field '\"' ||
 292       String.contains field '\n'
 293     then (
 294       let buffer = Buffer.create 100 in
 295       Buffer.add_char buffer '\"';
 296       for i = 0 to (String.length field) - 1 do
 297         match field.[i] with
 298             '\"' -> Buffer.add_string buffer "\"\""
 299           | c    -> Buffer.add_char buffer c
 300       done;
 301       Buffer.add_char buffer '\"';
 302       Buffer.contents buffer
 303     )
 304     else
 305       field
 306   in
 307
 308   let separator = String.make 1 separator in
 309   List.iter (fun line ->
 310                output_string chan (String.concat separator
 311                                      (List.map quote_field line));
 312                output_char chan '\n') csv
 313
 314 let print ?separator csv =
 315   save_out ?separator stdout csv; flush stdout
 316
 317 let save ?separator file csv =
 318   let chan = open_out file in
 319   save_out ?separator chan csv;
 320   close_out chan
 321
 322 let save_out_readable chan csv =
 323   (* Escape all the strings in the CSV file first. *)
 324   let csv = List.map (List.map String.escaped) csv in
 325
 326   let csv = square csv in
 327
 328   (* Find the width of each column. *)
 329   let widths =
 330     match csv with
 331       | [] -> []
 332       | r :: _ ->
 333           let n = List.length r in
 334           let lengths = List.map (List.map String.length) csv in
 335           let max2rows r1 r2 =
 336             let rp = List.combine r1 r2 in
 337             List.map (fun ((a : int), (b : int)) -> max a b) rp
 338           in
 339           let rec repeat x = function
 340             | 0 -> []
 341             | i -> x :: repeat x (i-1)
 342           in
 343           List.fold_left max2rows (repeat 0 n) lengths in
 344
 345   (* Print out each cell at the correct width. *)
 346   let rec repeat f = function
 347     | 0 -> ()
 348     | i -> f (); repeat f (i-1)
 349   in
 350   List.iter (
 351     fun row ->
 352       let row = List.combine widths row in
 353       List.iter (
 354         fun (width, cell) ->
 355           output_string chan cell;
 356           let n = String.length cell in
 357           repeat (fun () -> output_char chan ' ') (width - n + 1)
 358       ) row;
 359       output_char chan '\n'
 360   ) csv
 361
 362 let print_readable = save_out_readable stdout