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 if field == `\.` || strings.ContainsRune(field, w.Comma) || strings.ContainsAny(field, "\"\r\n") { 162 return true 163 } 164 165 r1, _ := utf8.DecodeRuneInString(field) 166 return unicode.IsSpace(r1) 167 } 168
View as plain text