...
Run Format

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

View as plain text