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