Home | History | Annotate | Download | only in csv
      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 output line 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 	if !validDelim(w.Comma) {
     42 		return errInvalidDelim
     43 	}
     44 
     45 	for n, field := range record {
     46 		if n > 0 {
     47 			if _, err := w.w.WriteRune(w.Comma); err != nil {
     48 				return err
     49 			}
     50 		}
     51 
     52 		// If we don't have to have a quoted field then just
     53 		// write out the field and continue to the next field.
     54 		if !w.fieldNeedsQuotes(field) {
     55 			if _, err := w.w.WriteString(field); err != nil {
     56 				return err
     57 			}
     58 			continue
     59 		}
     60 		if err := w.w.WriteByte('"'); err != nil {
     61 			return err
     62 		}
     63 
     64 		for _, r1 := range field {
     65 			var err error
     66 			switch r1 {
     67 			case '"':
     68 				_, err = w.w.WriteString(`""`)
     69 			case '\r':
     70 				if !w.UseCRLF {
     71 					err = w.w.WriteByte('\r')
     72 				}
     73 			case '\n':
     74 				if w.UseCRLF {
     75 					_, err = w.w.WriteString("\r\n")
     76 				} else {
     77 					err = w.w.WriteByte('\n')
     78 				}
     79 			default:
     80 				_, err = w.w.WriteRune(r1)
     81 			}
     82 			if err != nil {
     83 				return err
     84 			}
     85 		}
     86 
     87 		if err := w.w.WriteByte('"'); err != nil {
     88 			return err
     89 		}
     90 	}
     91 	var err error
     92 	if w.UseCRLF {
     93 		_, err = w.w.WriteString("\r\n")
     94 	} else {
     95 		err = w.w.WriteByte('\n')
     96 	}
     97 	return err
     98 }
     99 
    100 // Flush writes any buffered data to the underlying io.Writer.
    101 // To check if an error occurred during the Flush, call Error.
    102 func (w *Writer) Flush() {
    103 	w.w.Flush()
    104 }
    105 
    106 // Error reports any error that has occurred during a previous Write or Flush.
    107 func (w *Writer) Error() error {
    108 	_, err := w.w.Write(nil)
    109 	return err
    110 }
    111 
    112 // WriteAll writes multiple CSV records to w using Write and then calls Flush.
    113 func (w *Writer) WriteAll(records [][]string) error {
    114 	for _, record := range records {
    115 		err := w.Write(record)
    116 		if err != nil {
    117 			return err
    118 		}
    119 	}
    120 	return w.w.Flush()
    121 }
    122 
    123 // fieldNeedsQuotes reports whether our field must be enclosed in quotes.
    124 // Fields with a Comma, fields with a quote or newline, and
    125 // fields which start with a space must be enclosed in quotes.
    126 // We used to quote empty strings, but we do not anymore (as of Go 1.4).
    127 // The two representations should be equivalent, but Postgres distinguishes
    128 // quoted vs non-quoted empty string during database imports, and it has
    129 // an option to force the quoted behavior for non-quoted CSV but it has
    130 // no option to force the non-quoted behavior for quoted CSV, making
    131 // CSV with quoted empty strings strictly less useful.
    132 // Not quoting the empty string also makes this package match the behavior
    133 // of Microsoft Excel and Google Drive.
    134 // For Postgres, quote the data terminating string `\.`.
    135 func (w *Writer) fieldNeedsQuotes(field string) bool {
    136 	if field == "" {
    137 		return false
    138 	}
    139 	if field == `\.` || strings.ContainsRune(field, w.Comma) || strings.ContainsAny(field, "\"\r\n") {
    140 		return true
    141 	}
    142 
    143 	r1, _ := utf8.DecodeRuneInString(field)
    144 	return unicode.IsSpace(r1)
    145 }
    146