-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathencode.go
More file actions
216 lines (200 loc) · 4.65 KB
/
encode.go
File metadata and controls
216 lines (200 loc) · 4.65 KB
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
package csvstruct
import (
"encoding"
"encoding/csv"
"errors"
"fmt"
"io"
"reflect"
"sort"
)
var textMarshalerType = reflect.TypeOf(new(encoding.TextMarshaler)).Elem()
// Encoder encodes and writes CSV rows to an output stream.
type Encoder interface {
// EncodeNext encodes v into a CSV row and writes it to the Encoder's
// Writer.
//
// On the first call to EncodeNext, v's fields will be used to write the
// header row, then v's values will be written as the second row.
EncodeNext(v interface{}) error
// Opts specifies options to modify encoding behavior.
//
// It returns the Encoder, to support chaining.
Opts(EncodeOpts) Encoder
}
// EncodeOpts specifies options to modify encoding behavior.
type EncodeOpts struct {
SkipHeader bool // True to skip writing the header row
Comma rune // Field delimiter (set to ',' by default)
UseCRLF bool // True to use \r\n as the line terminator
}
type encoder struct {
w csv.Writer
hm map[string]int
opts EncodeOpts
}
// NewEncoder returns an encoder that writes to w.
func NewEncoder(w io.Writer) Encoder {
csvw := csv.NewWriter(w)
return &encoder{w: *csvw}
}
func (e *encoder) Opts(opts EncodeOpts) Encoder {
if opts.Comma != rune(0) {
e.w.Comma = opts.Comma
}
e.w.UseCRLF = opts.UseCRLF
e.opts = opts
return e
}
func (e *encoder) EncodeNext(v interface{}) error {
if v == nil {
return nil
}
switch reflect.ValueOf(v).Type().Kind() {
case reflect.Map:
return e.encodeMap(v)
case reflect.Struct:
return e.encodeStruct(v)
default:
return errors.New("must encode map or struct")
}
}
func (e *encoder) encodeMap(v interface{}) error {
if reflect.ValueOf(v).Type().Key().Kind() != reflect.String {
return errors.New("map key must be string")
}
m := v.(map[string]interface{})
if e.hm == nil {
e.hm = make(map[string]int)
headers := []string{}
for k := range m {
headers = append(headers, k)
}
sort.Strings(headers)
for i, h := range headers {
e.hm[h] = i
}
if len(e.hm) == 0 {
// First row was an empty map, so write nothing.
// This will result in an empty output no matter what is Encoded.
return nil
}
if !e.opts.SkipHeader {
if err := e.w.Write(headers); err != nil {
return err
}
}
}
row := make([]string, len(m))
add := false // Whether there has been a row to write in this call.
for h, i := range e.hm {
val, ok := m[h]
if !ok {
continue
}
add = true
row[i] = fmt.Sprint(val)
}
if !add {
return nil
}
if err := e.w.Write(row); err != nil {
return err
}
e.w.Flush()
return e.w.Error()
}
func (e *encoder) encodeStruct(v interface{}) error {
t := reflect.ValueOf(v).Type()
if e.hm == nil {
e.hm = make(map[string]int)
headers := []string{}
i := 0
for j := 0; j < t.NumField(); j++ {
f := t.Field(j)
if f.Anonymous {
continue
}
if f.PkgPath != "" { // Filter unexported fields
continue
}
n := f.Name
if f.Tag.Get("csv") != "" {
n = f.Tag.Get("csv")
if n == "-" {
continue
}
}
headers = append(headers, n)
e.hm[n] = i
i++
}
if len(e.hm) == 0 {
// Header row has no exported, unignored fields, so write nothing.
// This will result in an empty output no matter what is Encoded.
return nil
}
if !e.opts.SkipHeader {
if err := e.w.Write(headers); err != nil {
return err
}
}
}
rv := reflect.ValueOf(v)
row := make([]string, len(e.hm))
add := false // Whether there has been a row to write in this call.
for i := 0; i < t.NumField(); i++ {
f := t.Field(i)
if f.PkgPath != "" { // Filter unexported fields
continue
}
n := f.Name
if f.Tag.Get("csv") != "" {
n = f.Tag.Get("csv")
}
fi, ok := e.hm[n]
if !ok {
// Unmapped header value
continue
}
add = true
vf := rv.Field(i)
if vf.Type().Implements(textMarshalerType) {
if tm, ok := vf.Interface().(encoding.TextMarshaler); ok {
b, err := tm.MarshalText()
if err != nil {
return err
}
row[fi] = string(b)
continue
} else {
panic("unreachable")
}
}
if vf.Kind() == reflect.Ptr {
vf = vf.Elem()
}
switch vf.Kind() {
case reflect.String:
row[fi] = vf.String()
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
row[fi] = fmt.Sprintf("%d", vf.Int())
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
row[fi] = fmt.Sprintf("%d", vf.Uint())
case reflect.Float64:
row[fi] = fmt.Sprintf("%f", vf.Float())
case reflect.Bool:
row[fi] = fmt.Sprintf("%t", vf.Bool())
default:
return fmt.Errorf("can't encode type %v", f.Type)
}
}
if !add {
return nil
}
if err := e.w.Write(row); err != nil {
return err
}
e.w.Flush()
return e.w.Error()
}