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 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  
    61  		if err := w.w.WriteByte('"'); err != nil {
    62  			return err
    63  		}
    64  		for len(field) > 0 {
    65  			// Search for special characters.
    66  			i := strings.IndexAny(field, "\"\r\n")
    67  			if i < 0 {
    68  				i = len(field)
    69  			}
    70  
    71  			// Copy verbatim everything before the special character.
    72  			if _, err := w.w.WriteString(field[:i]); err != nil {
    73  				return err
    74  			}
    75  			field = field[i:]
    76  
    77  			// Encode the special character.
    78  			if len(field) > 0 {
    79  				var err error
    80  				switch field[0] {
    81  				case '"':
    82  					_, err = w.w.WriteString(`""`)
    83  				case '\r':
    84  					if !w.UseCRLF {
    85  						err = w.w.WriteByte('\r')
    86  					}
    87  				case '\n':
    88  					if w.UseCRLF {
    89  						_, err = w.w.WriteString("\r\n")
    90  					} else {
    91  						err = w.w.WriteByte('\n')
    92  					}
    93  				}
    94  				field = field[1:]
    95  				if err != nil {
    96  					return err
    97  				}
    98  			}
    99  		}
   100  		if err := w.w.WriteByte('"'); err != nil {
   101  			return err
   102  		}
   103  	}
   104  	var err error
   105  	if w.UseCRLF {
   106  		_, err = w.w.WriteString("\r\n")
   107  	} else {
   108  		err = w.w.WriteByte('\n')
   109  	}
   110  	return err
   111  }
   112  
   113  // Flush writes any buffered data to the underlying io.Writer.
   114  // To check if an error occurred during the Flush, call Error.
   115  func (w *Writer) Flush() {
   116  	w.w.Flush()
   117  }
   118  
   119  // Error reports any error that has occurred during a previous Write or Flush.
   120  func (w *Writer) Error() error {
   121  	_, err := w.w.Write(nil)
   122  	return err
   123  }
   124  
   125  // WriteAll writes multiple CSV records to w using Write and then calls Flush.
   126  func (w *Writer) WriteAll(records [][]string) error {
   127  	for _, record := range records {
   128  		err := w.Write(record)
   129  		if err != nil {
   130  			return err
   131  		}
   132  	}
   133  	return w.w.Flush()
   134  }
   135  
   136  // fieldNeedsQuotes reports whether our field must be enclosed in quotes.
   137  // Fields with a Comma, fields with a quote or newline, and
   138  // fields which start with a space must be enclosed in quotes.
   139  // We used to quote empty strings, but we do not anymore (as of Go 1.4).
   140  // The two representations should be equivalent, but Postgres distinguishes
   141  // quoted vs non-quoted empty string during database imports, and it has
   142  // an option to force the quoted behavior for non-quoted CSV but it has
   143  // no option to force the non-quoted behavior for quoted CSV, making
   144  // CSV with quoted empty strings strictly less useful.
   145  // Not quoting the empty string also makes this package match the behavior
   146  // of Microsoft Excel and Google Drive.
   147  // For Postgres, quote the data terminating string `\.`.
   148  func (w *Writer) fieldNeedsQuotes(field string) bool {
   149  	if field == "" {
   150  		return false
   151  	}
   152  	if field == `\.` || strings.ContainsRune(field, w.Comma) || strings.ContainsAny(field, "\"\r\n") {
   153  		return true
   154  	}
   155  
   156  	r1, _ := utf8.DecodeRuneInString(field)
   157  	return unicode.IsSpace(r1)
   158  }
   159
View as plain text