...

Source file src/pkg/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 using CSV encoding.
    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	//
    25	// The writes of individual records are buffered.
    26	// After all data has been written, the client should call the
    27	// Flush method to guarantee all data has been forwarded to
    28	// the underlying io.Writer.  Any errors that occurred should
    29	// be checked by calling the Error method.
    30	type Writer struct {
    31		Comma   rune // Field delimiter (set to ',' by NewWriter)
    32		UseCRLF bool // True to use \r\n as the line terminator
    33		w       *bufio.Writer
    34	}
    35	
    36	// NewWriter returns a new Writer that writes to w.
    37	func NewWriter(w io.Writer) *Writer {
    38		return &Writer{
    39			Comma: ',',
    40			w:     bufio.NewWriter(w),
    41		}
    42	}
    43	
    44	// Write writes a single CSV record to w along with any necessary quoting.
    45	// A record is a slice of strings with each string being one field.
    46	// Writes are buffered, so Flush must eventually be called to ensure
    47	// that the record is written to the underlying io.Writer.
    48	func (w *Writer) Write(record []string) error {
    49		if !validDelim(w.Comma) {
    50			return errInvalidDelim
    51		}
    52	
    53		for n, field := range record {
    54			if n > 0 {
    55				if _, err := w.w.WriteRune(w.Comma); err != nil {
    56					return err
    57				}
    58			}
    59	
    60			// If we don't have to have a quoted field then just
    61			// write out the field and continue to the next field.
    62			if !w.fieldNeedsQuotes(field) {
    63				if _, err := w.w.WriteString(field); err != nil {
    64					return err
    65				}
    66				continue
    67			}
    68	
    69			if err := w.w.WriteByte('"'); err != nil {
    70				return err
    71			}
    72			for len(field) > 0 {
    73				// Search for special characters.
    74				i := strings.IndexAny(field, "\"\r\n")
    75				if i < 0 {
    76					i = len(field)
    77				}
    78	
    79				// Copy verbatim everything before the special character.
    80				if _, err := w.w.WriteString(field[:i]); err != nil {
    81					return err
    82				}
    83				field = field[i:]
    84	
    85				// Encode the special character.
    86				if len(field) > 0 {
    87					var err error
    88					switch field[0] {
    89					case '"':
    90						_, err = w.w.WriteString(`""`)
    91					case '\r':
    92						if !w.UseCRLF {
    93							err = w.w.WriteByte('\r')
    94						}
    95					case '\n':
    96						if w.UseCRLF {
    97							_, err = w.w.WriteString("\r\n")
    98						} else {
    99							err = w.w.WriteByte('\n')
   100						}
   101					}
   102					field = field[1:]
   103					if err != nil {
   104						return err
   105					}
   106				}
   107			}
   108			if err := w.w.WriteByte('"'); err != nil {
   109				return err
   110			}
   111		}
   112		var err error
   113		if w.UseCRLF {
   114			_, err = w.w.WriteString("\r\n")
   115		} else {
   116			err = w.w.WriteByte('\n')
   117		}
   118		return err
   119	}
   120	
   121	// Flush writes any buffered data to the underlying io.Writer.
   122	// To check if an error occurred during the Flush, call Error.
   123	func (w *Writer) Flush() {
   124		w.w.Flush()
   125	}
   126	
   127	// Error reports any error that has occurred during a previous Write or Flush.
   128	func (w *Writer) Error() error {
   129		_, err := w.w.Write(nil)
   130		return err
   131	}
   132	
   133	// WriteAll writes multiple CSV records to w using Write and then calls Flush,
   134	// returning any error from the Flush.
   135	func (w *Writer) WriteAll(records [][]string) error {
   136		for _, record := range records {
   137			err := w.Write(record)
   138			if err != nil {
   139				return err
   140			}
   141		}
   142		return w.w.Flush()
   143	}
   144	
   145	// fieldNeedsQuotes reports whether our field must be enclosed in quotes.
   146	// Fields with a Comma, fields with a quote or newline, and
   147	// fields which start with a space must be enclosed in quotes.
   148	// We used to quote empty strings, but we do not anymore (as of Go 1.4).
   149	// The two representations should be equivalent, but Postgres distinguishes
   150	// quoted vs non-quoted empty string during database imports, and it has
   151	// an option to force the quoted behavior for non-quoted CSV but it has
   152	// no option to force the non-quoted behavior for quoted CSV, making
   153	// CSV with quoted empty strings strictly less useful.
   154	// Not quoting the empty string also makes this package match the behavior
   155	// of Microsoft Excel and Google Drive.
   156	// For Postgres, quote the data terminating string `\.`.
   157	func (w *Writer) fieldNeedsQuotes(field string) bool {
   158		if field == "" {
   159			return false
   160		}
   161		if field == `\.` || strings.ContainsRune(field, w.Comma) || strings.ContainsAny(field, "\"\r\n") {
   162			return true
   163		}
   164	
   165		r1, _ := utf8.DecodeRuneInString(field)
   166		return unicode.IsSpace(r1)
   167	}
   168	

View as plain text