libgo/go/encoding/csv/writer.go

   1 // Copyright 2011 The Go Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 package csv
   6
   7 import (
   8         "bufio"
   9         "io"
  10         "strings"
  11         "unicode"
  12         "unicode/utf8"
  13 )
  14
  15 // A Writer writes records to a CSV encoded file.
  16 //
  17 // As returned by NewWriter, a Writer writes records terminated by a
  18 // newline and uses ',' as the field delimiter. The exported fields can be
  19 // changed to customize the details before the first call to Write or WriteAll.
  20 //
  21 // Comma is the field delimiter.
  22 //
  23 // If UseCRLF is true, the Writer ends each record with \r\n instead of \n.
  24 type Writer struct {
  25         Comma   rune // Field delimiter (set to ',' by NewWriter)
  26         UseCRLF bool // True to use \r\n as the line terminator
  27         w       *bufio.Writer
  28 }
  29
  30 // NewWriter returns a new Writer that writes to w.
  31 func NewWriter(w io.Writer) *Writer {
  32         return &Writer{
  33                 Comma: ',',
  34                 w:     bufio.NewWriter(w),
  35         }
  36 }
  37
  38 // Writer writes a single CSV record to w along with any necessary quoting.
  39 // A record is a slice of strings with each string being one field.
  40 func (w *Writer) Write(record []string) error {
  41         for n, field := range record {
  42                 if n > 0 {
  43                         if _, err := w.w.WriteRune(w.Comma); err != nil {
  44                                 return err
  45                         }
  46                 }
  47
  48                 // If we don't have to have a quoted field then just
  49                 // write out the field and continue to the next field.
  50                 if !w.fieldNeedsQuotes(field) {
  51                         if _, err := w.w.WriteString(field); err != nil {
  52                                 return err
  53                         }
  54                         continue
  55                 }
  56                 if err := w.w.WriteByte('"'); err != nil {
  57                         return err
  58                 }
  59
  60                 for _, r1 := range field {
  61                         var err error
  62                         switch r1 {
  63                         case '"':
  64                                 _, err = w.w.WriteString(`""`)
  65                         case '\r':
  66                                 if !w.UseCRLF {
  67                                         err = w.w.WriteByte('\r')
  68                                 }
  69                         case '\n':
  70                                 if w.UseCRLF {
  71                                         _, err = w.w.WriteString("\r\n")
  72                                 } else {
  73                                         err = w.w.WriteByte('\n')
  74                                 }
  75                         default:
  76                                 _, err = w.w.WriteRune(r1)
  77                         }
  78                         if err != nil {
  79                                 return err
  80                         }
  81                 }
  82
  83                 if err := w.w.WriteByte('"'); err != nil {
  84                         return err
  85                 }
  86         }
  87         var err error
  88         if w.UseCRLF {
  89                 _, err = w.w.WriteString("\r\n")
  90         } else {
  91                 err = w.w.WriteByte('\n')
  92         }
  93         return err
  94 }
  95
  96 // Flush writes any buffered data to the underlying io.Writer.
  97 // To check if an error occurred during the Flush, call Error.
  98 func (w *Writer) Flush() {
  99         w.w.Flush()
 100 }
 101
 102 // Error reports any error that has occurred during a previous Write or Flush.
 103 func (w *Writer) Error() error {
 104         _, err := w.w.Write(nil)
 105         return err
 106 }
 107
 108 // WriteAll writes multiple CSV records to w using Write and then calls Flush.
 109 func (w *Writer) WriteAll(records [][]string) error {
 110         for _, record := range records {
 111                 err := w.Write(record)
 112                 if err != nil {
 113                         return err
 114                 }
 115         }
 116         return w.w.Flush()
 117 }
 118
 119 // fieldNeedsQuotes reports whether our field must be enclosed in quotes.
 120 // Fields with a Comma, fields with a quote or newline, and
 121 // fields which start with a space must be enclosed in quotes.
 122 // We used to quote empty strings, but we do not anymore (as of Go 1.4).
 123 // The two representations should be equivalent, but Postgres distinguishes
 124 // quoted vs non-quoted empty string during database imports, and it has
 125 // an option to force the quoted behavior for non-quoted CSV but it has
 126 // no option to force the non-quoted behavior for quoted CSV, making
 127 // CSV with quoted empty strings strictly less useful.
 128 // Not quoting the empty string also makes this package match the behavior
 129 // of Microsoft Excel and Google Drive.
 130 // For Postgres, quote the data terminating string `\.`.
 131 func (w *Writer) fieldNeedsQuotes(field string) bool {
 132         if field == "" {
 133                 return false
 134         }
 135         if field == `\.` || strings.ContainsRune(field, w.Comma) || strings.ContainsAny(field, "\"\r\n") {
 136                 return true
 137         }
 138
 139         r1, _ := utf8.DecodeRuneInString(field)
 140         return unicode.IsSpace(r1)
 141 }