/
csv.go
107 lines (87 loc) · 1.66 KB
/
csv.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
package exchange
import (
"encoding/csv"
"errors"
"io"
"strings"
)
func NewCSVReader(rc io.ReadCloser) (Reader, error) {
if rc == nil {
return nil, errors.New("reader is nil")
}
defer rc.Close()
r := csv.NewReader(rc)
r.TrimLeadingSpace = true
header, err := r.Read()
if err != nil {
return nil, err
}
header = trimStringSliceSpace(header)
records, err := r.ReadAll()
if err != nil {
return nil, err
}
for i, r := range records {
records[i] = trimStringSliceSpace(r)
}
return &csvReader{
header: header,
records: records,
curr: 0,
total: len(records),
}, nil
}
type csvReader struct {
header []string
records [][]string
curr int
total int
}
var _ Reader = (*csvReader)(nil)
func (c *csvReader) Header() []string {
return c.header
}
func (c *csvReader) ReadRow() ([]string, error) {
if c.total < c.curr {
return nil, errors.New("no more row")
}
return c.records[c.curr-1], nil
}
func (c *csvReader) Next() bool {
if c.total <= c.curr {
return false
}
c.curr++
return true
}
func (c *csvReader) Total() uint {
return uint(c.total)
}
func trimStringSliceSpace(rs []string) []string {
nrs := make([]string, 0, len(rs))
for _, r := range rs {
nrs = append(nrs, strings.TrimSpace(r))
}
return nrs
}
func NewCSVWriter(w io.Writer) (Writer, error) {
if w == nil {
return nil, errors.New("writer is nil")
}
return &csvWriter{
w: csv.NewWriter(w),
}, nil
}
type csvWriter struct {
w *csv.Writer
}
func (c *csvWriter) WriteHeader(h []string) error {
return c.w.Write(h)
}
func (c *csvWriter) WriteRow(r []string) error {
return c.w.Write(r)
}
func (c *csvWriter) Flush() error {
c.w.Flush()
return nil
}