csv.go 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319
  1. // Copyright 2014 Jonathan Picques. All rights reserved.
  2. // Use of this source code is governed by a MIT license
  3. // The license can be found in the LICENSE file.
  4. // The GoCSV package aims to provide easy CSV serialization and deserialization to the golang programming language
  5. package gocsv
  6. import (
  7. "bytes"
  8. "encoding/csv"
  9. "fmt"
  10. "io"
  11. "os"
  12. "reflect"
  13. "strings"
  14. )
  15. // FailIfUnmatchedStructTags indicates whether it is considered an error when there is an unmatched
  16. // struct tag.
  17. var FailIfUnmatchedStructTags = false
  18. // FailIfDoubleHeaderNames indicates whether it is considered an error when a header name is repeated
  19. // in the csv header.
  20. var FailIfDoubleHeaderNames = false
  21. // ShouldAlignDuplicateHeadersWithStructFieldOrder indicates whether we should align duplicate CSV
  22. // headers per their alignment in the struct definition.
  23. var ShouldAlignDuplicateHeadersWithStructFieldOrder = false
  24. // TagSeparator defines seperator string for multiple csv tags in struct fields
  25. var TagSeparator = ","
  26. // --------------------------------------------------------------------------
  27. // CSVWriter used to format CSV
  28. var selfCSVWriter = DefaultCSVWriter
  29. // DefaultCSVWriter is the default SafeCSVWriter used to format CSV (cf. csv.NewWriter)
  30. func DefaultCSVWriter(out io.Writer) *SafeCSVWriter {
  31. writer := NewSafeCSVWriter(csv.NewWriter(out))
  32. // As only one rune can be defined as a CSV separator, we are going to trim
  33. // the custom tag separator and use the first rune.
  34. if runes := []rune(strings.TrimSpace(TagSeparator)); len(runes) > 0 {
  35. writer.Comma = runes[0]
  36. }
  37. return writer
  38. }
  39. // SetCSVWriter sets the SafeCSVWriter used to format CSV.
  40. func SetCSVWriter(csvWriter func(io.Writer) *SafeCSVWriter) {
  41. selfCSVWriter = csvWriter
  42. }
  43. func getCSVWriter(out io.Writer) *SafeCSVWriter {
  44. return selfCSVWriter(out)
  45. }
  46. // --------------------------------------------------------------------------
  47. // CSVReader used to parse CSV
  48. var selfCSVReader = DefaultCSVReader
  49. // DefaultCSVReader is the default CSV reader used to parse CSV (cf. csv.NewReader)
  50. func DefaultCSVReader(in io.Reader) CSVReader {
  51. return csv.NewReader(in)
  52. }
  53. // LazyCSVReader returns a lazy CSV reader, with LazyQuotes and TrimLeadingSpace.
  54. func LazyCSVReader(in io.Reader) CSVReader {
  55. csvReader := csv.NewReader(in)
  56. csvReader.LazyQuotes = true
  57. csvReader.TrimLeadingSpace = true
  58. return csvReader
  59. }
  60. // SetCSVReader sets the CSV reader used to parse CSV.
  61. func SetCSVReader(csvReader func(io.Reader) CSVReader) {
  62. selfCSVReader = csvReader
  63. }
  64. func getCSVReader(in io.Reader) CSVReader {
  65. return selfCSVReader(in)
  66. }
  67. // --------------------------------------------------------------------------
  68. // Marshal functions
  69. // MarshalFile saves the interface as CSV in the file.
  70. func MarshalFile(in interface{}, file *os.File) (err error) {
  71. return Marshal(in, file)
  72. }
  73. // MarshalString returns the CSV string from the interface.
  74. func MarshalString(in interface{}) (out string, err error) {
  75. bufferString := bytes.NewBufferString(out)
  76. if err := Marshal(in, bufferString); err != nil {
  77. return "", err
  78. }
  79. return bufferString.String(), nil
  80. }
  81. // MarshalBytes returns the CSV bytes from the interface.
  82. func MarshalBytes(in interface{}) (out []byte, err error) {
  83. bufferString := bytes.NewBuffer(out)
  84. if err := Marshal(in, bufferString); err != nil {
  85. return nil, err
  86. }
  87. return bufferString.Bytes(), nil
  88. }
  89. // Marshal returns the CSV in writer from the interface.
  90. func Marshal(in interface{}, out io.Writer) (err error) {
  91. writer := getCSVWriter(out)
  92. return writeTo(writer, in, false)
  93. }
  94. // Marshal returns the CSV in writer from the interface.
  95. func MarshalWithoutHeaders(in interface{}, out io.Writer) (err error) {
  96. writer := getCSVWriter(out)
  97. return writeTo(writer, in, true)
  98. }
  99. // MarshalChan returns the CSV read from the channel.
  100. func MarshalChan(c <-chan interface{}, out *SafeCSVWriter) error {
  101. return writeFromChan(out, c)
  102. }
  103. // MarshalCSV returns the CSV in writer from the interface.
  104. func MarshalCSV(in interface{}, out *SafeCSVWriter) (err error) {
  105. return writeTo(out, in, false)
  106. }
  107. // MarshalCSVWithoutHeaders returns the CSV in writer from the interface.
  108. func MarshalCSVWithoutHeaders(in interface{}, out *SafeCSVWriter) (err error) {
  109. return writeTo(out, in, true)
  110. }
  111. // --------------------------------------------------------------------------
  112. // Unmarshal functions
  113. // UnmarshalFile parses the CSV from the file in the interface.
  114. func UnmarshalFile(in *os.File, out interface{}) error {
  115. return Unmarshal(in, out)
  116. }
  117. // UnmarshalString parses the CSV from the string in the interface.
  118. func UnmarshalString(in string, out interface{}) error {
  119. return Unmarshal(strings.NewReader(in), out)
  120. }
  121. // UnmarshalBytes parses the CSV from the bytes in the interface.
  122. func UnmarshalBytes(in []byte, out interface{}) error {
  123. return Unmarshal(bytes.NewReader(in), out)
  124. }
  125. // Unmarshal parses the CSV from the reader in the interface.
  126. func Unmarshal(in io.Reader, out interface{}) error {
  127. return readTo(newDecoder(in), out)
  128. }
  129. // UnmarshalDecoder parses the CSV from the decoder in the interface
  130. func UnmarshalDecoder(in Decoder, out interface{}) error {
  131. return readTo(in, out)
  132. }
  133. // UnmarshalCSV parses the CSV from the reader in the interface.
  134. func UnmarshalCSV(in CSVReader, out interface{}) error {
  135. return readTo(csvDecoder{in}, out)
  136. }
  137. // UnmarshalToChan parses the CSV from the reader and send each value in the chan c.
  138. // The channel must have a concrete type.
  139. func UnmarshalToChan(in io.Reader, c interface{}) error {
  140. if c == nil {
  141. return fmt.Errorf("goscv: channel is %v", c)
  142. }
  143. return readEach(newDecoder(in), c)
  144. }
  145. // UnmarshalDecoderToChan parses the CSV from the decoder and send each value in the chan c.
  146. // The channel must have a concrete type.
  147. func UnmarshalDecoderToChan(in SimpleDecoder, c interface{}) error {
  148. if c == nil {
  149. return fmt.Errorf("goscv: channel is %v", c)
  150. }
  151. return readEach(in, c)
  152. }
  153. // UnmarshalStringToChan parses the CSV from the string and send each value in the chan c.
  154. // The channel must have a concrete type.
  155. func UnmarshalStringToChan(in string, c interface{}) error {
  156. return UnmarshalToChan(strings.NewReader(in), c)
  157. }
  158. // UnmarshalBytesToChan parses the CSV from the bytes and send each value in the chan c.
  159. // The channel must have a concrete type.
  160. func UnmarshalBytesToChan(in []byte, c interface{}) error {
  161. return UnmarshalToChan(bytes.NewReader(in), c)
  162. }
  163. // UnmarshalToCallback parses the CSV from the reader and send each value to the given func f.
  164. // The func must look like func(Struct).
  165. func UnmarshalToCallback(in io.Reader, f interface{}) error {
  166. valueFunc := reflect.ValueOf(f)
  167. t := reflect.TypeOf(f)
  168. if t.NumIn() != 1 {
  169. return fmt.Errorf("the given function must have exactly one parameter")
  170. }
  171. cerr := make(chan error)
  172. c := reflect.MakeChan(reflect.ChanOf(reflect.BothDir, t.In(0)), 0)
  173. go func() {
  174. cerr <- UnmarshalToChan(in, c.Interface())
  175. }()
  176. for {
  177. select {
  178. case err := <-cerr:
  179. return err
  180. default:
  181. }
  182. v, notClosed := c.Recv()
  183. if !notClosed || v.Interface() == nil {
  184. break
  185. }
  186. valueFunc.Call([]reflect.Value{v})
  187. }
  188. return nil
  189. }
  190. // UnmarshalDecoderToCallback parses the CSV from the decoder and send each value to the given func f.
  191. // The func must look like func(Struct).
  192. func UnmarshalDecoderToCallback(in SimpleDecoder, f interface{}) error {
  193. valueFunc := reflect.ValueOf(f)
  194. t := reflect.TypeOf(f)
  195. if t.NumIn() != 1 {
  196. return fmt.Errorf("the given function must have exactly one parameter")
  197. }
  198. cerr := make(chan error)
  199. c := reflect.MakeChan(reflect.ChanOf(reflect.BothDir, t.In(0)), 0)
  200. go func() {
  201. cerr <- UnmarshalDecoderToChan(in, c.Interface())
  202. }()
  203. for {
  204. select {
  205. case err := <-cerr:
  206. return err
  207. default:
  208. }
  209. v, notClosed := c.Recv()
  210. if !notClosed || v.Interface() == nil {
  211. break
  212. }
  213. valueFunc.Call([]reflect.Value{v})
  214. }
  215. return nil
  216. }
  217. // UnmarshalBytesToCallback parses the CSV from the bytes and send each value to the given func f.
  218. // The func must look like func(Struct).
  219. func UnmarshalBytesToCallback(in []byte, f interface{}) error {
  220. return UnmarshalToCallback(bytes.NewReader(in), f)
  221. }
  222. // UnmarshalStringToCallback parses the CSV from the string and send each value to the given func f.
  223. // The func must look like func(Struct).
  224. func UnmarshalStringToCallback(in string, c interface{}) (err error) {
  225. return UnmarshalToCallback(strings.NewReader(in), c)
  226. }
  227. // CSVToMap creates a simple map from a CSV of 2 columns.
  228. func CSVToMap(in io.Reader) (map[string]string, error) {
  229. decoder := newDecoder(in)
  230. header, err := decoder.getCSVRow()
  231. if err != nil {
  232. return nil, err
  233. }
  234. if len(header) != 2 {
  235. return nil, fmt.Errorf("maps can only be created for csv of two columns")
  236. }
  237. m := make(map[string]string)
  238. for {
  239. line, err := decoder.getCSVRow()
  240. if err == io.EOF {
  241. break
  242. } else if err != nil {
  243. return nil, err
  244. }
  245. m[line[0]] = line[1]
  246. }
  247. return m, nil
  248. }
  249. // CSVToMaps takes a reader and returns an array of dictionaries, using the header row as the keys
  250. func CSVToMaps(reader io.Reader) ([]map[string]string, error) {
  251. r := csv.NewReader(reader)
  252. rows := []map[string]string{}
  253. var header []string
  254. for {
  255. record, err := r.Read()
  256. if err == io.EOF {
  257. break
  258. }
  259. if err != nil {
  260. return nil, err
  261. }
  262. if header == nil {
  263. header = record
  264. } else {
  265. dict := map[string]string{}
  266. for i := range header {
  267. dict[header[i]] = record[i]
  268. }
  269. rows = append(rows, dict)
  270. }
  271. }
  272. return rows, nil
  273. }