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  
   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