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 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 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) (err error) {
    41		for n, field := range record {
    42			if n > 0 {
    43				if _, err = w.w.WriteRune(w.Comma); err != nil {
    44					return
    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
    53				}
    54				continue
    55			}
    56			if err = w.w.WriteByte('"'); err != nil {
    57				return
    58			}
    59	
    60			for _, r1 := range field {
    61				switch r1 {
    62				case '"':
    63					_, err = w.w.WriteString(`""`)
    64				case '\r':
    65					if !w.UseCRLF {
    66						err = w.w.WriteByte('\r')
    67					}
    68				case '\n':
    69					if w.UseCRLF {
    70						_, err = w.w.WriteString("\r\n")
    71					} else {
    72						err = w.w.WriteByte('\n')
    73					}
    74				default:
    75					_, err = w.w.WriteRune(r1)
    76				}
    77				if err != nil {
    78					return
    79				}
    80			}
    81	
    82			if err = w.w.WriteByte('"'); err != nil {
    83				return
    84			}
    85		}
    86		if w.UseCRLF {
    87			_, err = w.w.WriteString("\r\n")
    88		} else {
    89			err = w.w.WriteByte('\n')
    90		}
    91		return
    92	}
    93	
    94	// Flush writes any buffered data to the underlying io.Writer.
    95	func (w *Writer) Flush() {
    96		w.w.Flush()
    97	}
    98	
    99	// WriteAll writes multiple CSV records to w using Write and then calls Flush.
   100	func (w *Writer) WriteAll(records [][]string) (err error) {
   101		for _, record := range records {
   102			err = w.Write(record)
   103			if err != nil {
   104				break
   105			}
   106		}
   107		w.Flush()
   108		return nil
   109	}
   110	
   111	// fieldNeedsQuotes returns true if our field must be enclosed in quotes.
   112	// Empty fields, files with a Comma, fields with a quote or newline, and
   113	// fields which start with a space must be enclosed in quotes.
   114	func (w *Writer) fieldNeedsQuotes(field string) bool {
   115		if len(field) == 0 || strings.IndexRune(field, w.Comma) >= 0 || strings.IndexAny(field, "\"\r\n") >= 0 {
   116			return true
   117		}
   118	
   119		r1, _ := utf8.DecodeRuneInString(field)
   120		return unicode.IsSpace(r1)
   121	}