basis/csv/csv.factor

   1 ! Copyright (C) 2007, 2008 Phil Dawes
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: kernel sequences io namespaces make combinators
   4 unicode.categories io.files combinators.short-circuit
   5 io.streams.string fry memoize ;
   6 IN: csv
   7
   8 SYMBOL: delimiter
   9
  10 CHAR: , delimiter set-global
  11
  12 <PRIVATE
  13
  14 : delimiter> ( -- delimiter ) delimiter get ; inline
  15
  16 MEMO: (field-end) ( delimiter -- delimiter' )
  17     "\n" swap suffix ; inline
  18
  19 : skip-to-field-end ( -- endchar )
  20     delimiter> (field-end) read-until nip ; inline
  21
  22 DEFER: quoted-field
  23
  24 MEMO: (quoted-field) ( delimiter -- delimiter' )
  25     "\"\n" swap suffix ; inline
  26
  27 : not-quoted-field ( -- endchar )
  28     delimiter> (quoted-field) read-until
  29     dup CHAR: " =
  30     [ 2drop quoted-field ]
  31     [ swap [ blank? ] trim % ]
  32     if ;
  33
  34 : maybe-escaped-quote ( -- endchar )
  35     read1 dup {
  36         { CHAR: "    [ , quoted-field ] }
  37         { delimiter> [ ] }
  38         { CHAR: \n   [ ] }
  39         [ 2drop skip-to-field-end ]
  40     } case ;
  41
  42 : quoted-field ( -- endchar )
  43     "\"" read-until
  44     drop % maybe-escaped-quote ;
  45
  46 : field ( -- sep string )
  47     [ not-quoted-field ] "" make  ;
  48
  49 : (row) ( -- sep )
  50     f delimiter> '[ dup _ = ]
  51     [ drop field , ] do while ;
  52
  53 : row ( -- eof? array[string] )
  54     [ (row) ] { } make ;
  55
  56 : (csv) ( -- )
  57     [ dup [ empty? ] all? [ drop ] [ , ] if ]
  58     [ row ] do while ;
  59
  60 PRIVATE>
  61
  62 : csv-row ( stream -- row )
  63     [ row nip ] with-input-stream ;
  64
  65 : csv ( stream -- rows )
  66     [ [ (csv) ] { } make ] with-input-stream
  67     dup last { "" } = [ but-last ] when ;
  68
  69 : string>csv ( string -- csv )
  70     <string-reader> csv ;
  71
  72 : file>csv ( path encoding -- csv )
  73     <file-reader> csv ;
  74
  75 : with-delimiter ( ch quot -- )
  76     [ delimiter ] dip with-variable ; inline
  77
  78 <PRIVATE
  79
  80 : needs-escaping? ( cell -- ? )
  81     delimiter> '[
  82         dup "\n\"" member? [ drop t ] [ _ = ] if
  83     ] any? ; inline
  84
  85 : escape-quotes ( cell -- cell' )
  86     [
  87         [
  88             [ , ]
  89             [ dup CHAR: " = [ , ] [ drop ] if ] bi
  90         ] each
  91     ] "" make ; inline
  92
  93 : enclose-in-quotes ( cell -- cell' )
  94     "\"" dup surround ; inline
  95
  96 : escape-if-required ( cell -- cell' )
  97     dup needs-escaping?
  98     [ escape-quotes enclose-in-quotes ] when ; inline
  99
 100 PRIVATE>
 101
 102 : write-row ( row -- )
 103     delimiter> '[ _ write1 ]
 104     [ escape-if-required write ] interleave nl ; inline
 105
 106 <PRIVATE
 107
 108 : (write-csv) ( rows -- )
 109     [ write-row ] each ;
 110
 111 PRIVATE>
 112
 113 : write-csv ( rows stream -- )
 114     [ (write-csv) ] with-output-stream ;
 115
 116 : csv>string ( csv -- string )
 117     [ (write-csv) ] with-string-writer ;
 118
 119 : csv>file ( rows path encoding -- ) <file-writer> write-csv ;