From cbf5fccb69f5a1eb60c3e1b04fad58ee5d6c99b1 Mon Sep 17 00:00:00 2001
From: Phil Dawes <phil@phildawes.net>
Date: Fri, 25 Jul 2008 22:02:07 +0100
Subject: [PATCH] Added write-csv word

---
 extra/csv/csv-docs.factor  | 11 +++++++++--
 extra/csv/csv-tests.factor | 10 +++++++++-
 extra/csv/csv.factor       | 25 +++++++++++++++++++++++--
 3 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/extra/csv/csv-docs.factor b/extra/csv/csv-docs.factor
index c9f39900ab..e4741f4810 100644
--- a/extra/csv/csv-docs.factor
+++ b/extra/csv/csv-docs.factor
@@ -2,20 +2,27 @@ USING: help.syntax help.markup kernel prettyprint sequences ;
 IN: csv
 
 HELP: csv
-{ $values { "stream" "a stream" }
+{ $values { "stream" "an input stream" }
           { "rows" "an array of arrays of fields" } } 
 { $description "parses a csv stream into an array of row arrays"
 } ;
 
 HELP: csv-row
-{ $values { "stream" "a stream" }
+{ $values { "stream" "an input stream" }
           { "row" "an array of fields" } } 
 { $description "parses a row from a csv stream"
 } ;
 
+HELP: write-csv
+{ $values { "rows" "an sequence of sequences of strings" }
+          { "stream" "an output stream" } } 
+{ $description "writes csv to the output stream, escaping where necessary"
+} ;
+
 
 HELP: with-delimiter
 { $values { "char" "field delimiter (e.g. CHAR: \t)" }
           { "quot" "a quotation" } }
 { $description "Sets the field delimiter for csv or csv-row words "
 } ;
+
diff --git a/extra/csv/csv-tests.factor b/extra/csv/csv-tests.factor
index 7e96dbc0a6..8261ae104a 100644
--- a/extra/csv/csv-tests.factor
+++ b/extra/csv/csv-tests.factor
@@ -1,5 +1,5 @@
-USING: io.streams.string csv tools.test shuffle ;
 IN: csv.tests
+USING: io.streams.string csv tools.test shuffle kernel strings ;
 
 ! I like to name my unit tests
 : named-unit-test ( name output input -- ) 
@@ -68,3 +68,11 @@ IN: csv.tests
 [ { { "foo" "bar" }
     { "1"   "2" } } ]
 [ "foo,\"bar\"\n1,2" <string-reader> csv ] named-unit-test
+
+"can write csv too!"
+[ "foo1,bar1\nfoo2,bar2\n" ]
+[ { { "foo1" "bar1" } { "foo2" "bar2" } } <string-writer> tuck write-csv >string ] named-unit-test
+
+"escapes quotes commas and newlines when writing"
+[ "\"fo\"\"o1\",bar1\n\"fo\no2\",\"b,ar2\"\n" ]
+[ { { "fo\"o1" "bar1" } { "fo\no2" "b,ar2" } } <string-writer> tuck write-csv >string ] named-unit-test ! "
diff --git a/extra/csv/csv.factor b/extra/csv/csv.factor
index 8ba0832b29..3d1fb64492 100644
--- a/extra/csv/csv.factor
+++ b/extra/csv/csv.factor
@@ -10,7 +10,7 @@ IN: csv
 DEFER: quoted-field
 
 VAR: delimiter
-
+    
 ! trims whitespace from either end of string
 : trim-whitespace ( str -- str )
   [ blank? ] trim ; inline
@@ -57,7 +57,7 @@ VAR: delimiter
   [ (csv) ] when ;
 
 : init-vars ( -- )
-  delimiter> [ CHAR: , >delimiter ] unless ; inline
+  delimiter> [ CHAR: , >delimiter ] unless ; inline 
   
 : csv-row ( stream -- row )
   init-vars
@@ -69,3 +69,24 @@ VAR: delimiter
 
 : with-delimiter ( char quot -- )
   delimiter swap with-variable ; inline
+
+
+    
+: needs-escaping? ( cell -- ? )
+  [ "\n\"" delimiter> suffix member? ] contains? ; inline ! "
+
+: escape-quotes ( cell -- cell' )
+  [ [ dup , CHAR: " = [ CHAR: " , ] when ] each ] "" make ; inline
+
+: enclose-in-quotes ( cell -- cell' )
+  CHAR: " [ prefix ] [ suffix ] bi ; inline ! "
+    
+: escape-if-required ( cell -- cell' )
+  dup needs-escaping? [ escape-quotes enclose-in-quotes ] when ; inline
+    
+: write-row ( row -- )
+  [ delimiter> write1 ] [ escape-if-required write ] interleave nl ; inline
+    
+: write-csv ( rows outstream -- )
+  init-vars
+  [ [ write-row ] each ] with-output-stream ;