/
csvutil.go
162 lines (138 loc) · 3.86 KB
/
csvutil.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
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
package csvutil
import (
"encoding/csv"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"reflect"
"strconv"
)
// UnmarshalFile reads the CSV file from filepath and unmarshals it into v.
func UnmarshalFile[T any](filepath string) ([]T, error) {
f, err := os.Open(filepath)
if err != nil {
return nil, fmt.Errorf("failed to open %q: %w", filepath, err)
}
defer f.Close()
return Unmarshal[T](csv.NewReader(f))
}
// Unmarshal reads the CSV file from r and unmarshals it into v.
func Unmarshal[T any](r *csv.Reader) ([]T, error) {
var z T
zt := reflect.TypeOf(z)
if zt.Kind() != reflect.Struct {
panic("expected struct")
}
numFields := zt.NumField()
values := reflect.MakeSlice(reflect.SliceOf(zt), 0, 0)
value := reflect.New(zt).Elem()
for {
record, err := r.Read()
if err != nil {
if errors.Is(err, io.EOF) {
break
}
return nil, fmt.Errorf("failed to read CSV record: %w", err)
}
// Ensure at least numFields fields are present.
if len(record) < numFields {
return nil, fmt.Errorf("expected %d fields, got %d", numFields, len(record))
}
for i := 0; i < numFields; i++ {
if err := unmarshalCell(record[i], value.Field(i)); err != nil {
return nil, fmt.Errorf("failed to unmarshal field %d: %w", i, err)
}
}
values = reflect.Append(values, value)
}
return values.Interface().([]T), nil
}
func unmarshalCell(v string, dst reflect.Value) error {
t := dst.Type()
switch t.Kind() {
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
i, err := strconv.ParseInt(v, 10, t.Bits())
if err != nil {
return fmt.Errorf("cannot parse %q as %s: %w", v, t, err)
}
dst.SetInt(i)
return nil
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
i, err := strconv.ParseUint(v, 10, t.Bits())
if err != nil {
return fmt.Errorf("cannot parse %q as %s: %w", v, t, err)
}
dst.SetUint(i)
return nil
case reflect.Float32, reflect.Float64:
f, err := strconv.ParseFloat(v, t.Bits())
if err != nil {
return fmt.Errorf("cannot parse %q as %s: %w", v, t, err)
}
dst.SetFloat(f)
return nil
case reflect.String:
dst.SetString(v)
return nil
default:
return fmt.Errorf("cannot parse %q as %s: unsupported type", v, t)
}
}
// MarshalFile writes the CSV representation of values to filepath.
func MarshalFile[T any](filepath string, values []T) error {
f, err := os.Create(filepath)
if err != nil {
return fmt.Errorf("failed to create %q: %w", filepath, err)
}
if err := Marshal(csv.NewWriter(f), values); err != nil {
return fmt.Errorf("failed to marshal: %w", err)
}
return nil
}
// Marshal writes the CSV representation of values to w.
func Marshal[T any](w *csv.Writer, values []T) error {
var z T
zt := reflect.TypeOf(z)
if zt.Kind() != reflect.Struct {
panic("expected struct")
}
numFields := zt.NumField()
rvalues := reflect.ValueOf(values)
for i := range values {
rvalue := rvalues.Index(i)
record := make([]string, numFields)
for i := 0; i < numFields; i++ {
v, err := marshalField(rvalue, i)
if err != nil {
return fmt.Errorf(
"failed to marshal %s.%s: %w",
zt, zt.Field(i).Name, err)
}
record[i] = v
}
if err := w.Write(record); err != nil {
return fmt.Errorf("failed to write CSV record: %w", err)
}
}
w.Flush()
return w.Error()
}
func marshalField(rvalue reflect.Value, i int) (string, error) {
rfield := rvalue.Field(i)
switch rfield.Kind() {
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
reflect.Float32, reflect.Float64:
b, err := json.Marshal(rfield.Interface())
if err != nil {
return "", fmt.Errorf("failed to marshal as JSON: %w", err)
}
return string(b), nil
case reflect.String:
return rfield.String(), nil
default:
return "", fmt.Errorf("unsupported type %s", rfield.Type())
}
}