writer.go

Documentation: encoding/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 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  
	 162  	if field == `\.` {
	 163  		return true
	 164  	}
	 165  
	 166  	if w.Comma < utf8.RuneSelf {
	 167  		for i := 0; i < len(field); i++ {
	 168  			c := field[i]
	 169  			if c == '\n' || c == '\r' || c == '"' || c == byte(w.Comma) {
	 170  				return true
	 171  			}
	 172  		}
	 173  	} else {
	 174  		if strings.ContainsRune(field, w.Comma) || strings.ContainsAny(field, "\"\r\n") {
	 175  			return true
	 176  		}
	 177  	}
	 178  
	 179  	r1, _ := utf8.DecodeRuneInString(field)
	 180  	return unicode.IsSpace(r1)
	 181  }
	 182
View as plain text