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 ;
  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 ;
  26
  27 : not-quoted-field ( -- endchar )
  28     delimiter> (quoted-field) read-until
  29     dup {
  30         { CHAR: "    [ 2drop quoted-field ] }
  31         { delimiter> [ swap [ blank? ] trim % ] }
  32         { CHAR: \n   [ swap [ blank? ] trim % ] }
  33         { f          [ swap [ blank? ] trim % ] }
  34     } case ;
  35
  36 : maybe-escaped-quote ( -- endchar )
  37     read1 dup {
  38         { CHAR: "    [ , quoted-field ] }
  39         { delimiter> [ ] }
  40         { CHAR: \n   [ ] }
  41         [ 2drop skip-to-field-end ]
  42     } case ;
  43
  44 : quoted-field ( -- endchar )
  45     "\"" read-until
  46     drop % maybe-escaped-quote ;
  47
  48 : field ( -- sep string )
  49     [ not-quoted-field ] "" make  ;
  50
  51 : (row) ( -- sep )
  52     f delimiter> '[ dup _ = ]
  53     [ drop field , ] do while ;
  54
  55 : row ( -- eof? array[string] )
  56     [ (row) ] { } make ;
  57
  58 : (csv) ( -- )
  59     [ dup [ empty? ] all? [ drop ] [ , ] if ]
  60     [ row ] do while ;
  61
  62 PRIVATE>
  63
  64 : csv-row ( stream -- row )
  65     [ row nip ] with-input-stream ;
  66
  67 : csv ( stream -- rows )
  68     [ [ (csv) ] { } make ] with-input-stream
  69     dup last { "" } = [ but-last ] when ;
  70
  71 : string>csv ( string -- csv )
  72     <string-reader> csv ;
  73
  74 : file>csv ( path encoding -- csv )
  75     <file-reader> csv ;
  76
  77 : with-delimiter ( ch quot -- )
  78     [ delimiter ] dip with-variable ; inline
  79
  80 <PRIVATE
  81
  82 : needs-escaping? ( cell -- ? )
  83     delimiter> '[
  84         dup "\n\"" member? [ drop t ] [ _ = ] if
  85     ] any? ; inline
  86
  87 : escape-quotes ( cell -- cell' )
  88     [
  89         [
  90             [ , ]
  91             [ dup CHAR: " = [ , ] [ drop ] if ] bi
  92         ] each
  93     ] "" make ; inline
  94
  95 : enclose-in-quotes ( cell -- cell' )
  96     "\"" dup surround ; inline
  97
  98 : escape-if-required ( cell -- cell' )
  99     dup needs-escaping?
 100     [ escape-quotes enclose-in-quotes ] when ; inline
 101
 102 PRIVATE>
 103
 104 : write-row ( row -- )
 105     delimiter> '[ _ write1 ]
 106     [ escape-if-required write ] interleave nl ; inline
 107
 108 <PRIVATE
 109
 110 : (write-csv) ( rows -- )
 111     [ write-row ] each ;
 112
 113 PRIVATE>
 114
 115 : write-csv ( rows stream -- )
 116     [ (write-csv) ] with-output-stream ;
 117
 118 : csv>string ( csv -- string )
 119     [ (write-csv) ] with-string-writer ;
 120
 121 : csv>file ( rows path encoding -- ) <file-writer> write-csv ;