121 lines
1.9 KiB
Go
121 lines
1.9 KiB
Go
package builtins
|
|
|
|
import (
|
|
"context"
|
|
"encoding/csv"
|
|
"errors"
|
|
"io"
|
|
"io/fs"
|
|
"strings"
|
|
"ucl.lmika.dev/ucl"
|
|
)
|
|
|
|
type csvHandlers struct {
|
|
fs fs.FS
|
|
}
|
|
|
|
func CSV(fs fs.FS) ucl.Module {
|
|
fsh := csvHandlers{fs: fs}
|
|
|
|
return ucl.Module{
|
|
Name: "csv",
|
|
Builtins: map[string]ucl.BuiltinHandler{
|
|
"each-record": fsh.eachRecord,
|
|
},
|
|
}
|
|
}
|
|
|
|
func (h csvHandlers) eachRecord(ctx context.Context, args ucl.CallArgs) (any, error) {
|
|
var (
|
|
filename string
|
|
closure ucl.Invokable
|
|
)
|
|
|
|
if err := args.Bind(&filename, &closure); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
f, err := h.fs.Open(filename)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer f.Close()
|
|
|
|
cr := csv.NewReader(f)
|
|
|
|
header, err := cr.Read()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
hio := make(headerIndexObject)
|
|
for i, h := range header {
|
|
hio[h] = i
|
|
}
|
|
|
|
for {
|
|
record, err := cr.Read()
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
} else if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if _, err := closure.Invoke(ctx, stringSlice(record), hio); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
return nil, nil
|
|
}
|
|
|
|
type stringSlice []string
|
|
|
|
func (ss stringSlice) String() string {
|
|
return strings.Join(ss, ",")
|
|
}
|
|
func (ss stringSlice) Truthy() bool {
|
|
return len(ss) > 0
|
|
}
|
|
|
|
func (ss stringSlice) Len() int {
|
|
return len(ss)
|
|
}
|
|
func (ss stringSlice) Index(i int) ucl.Object {
|
|
return ucl.StringObject(ss[i])
|
|
}
|
|
|
|
type headerIndexObject map[string]int
|
|
|
|
func (hio headerIndexObject) String() string {
|
|
strs := make([]string, len(hio))
|
|
for h, i := range hio {
|
|
strs[i] = h
|
|
}
|
|
return strings.Join(strs, ",")
|
|
}
|
|
|
|
func (hio headerIndexObject) Truthy() bool {
|
|
return len(hio) > 0
|
|
}
|
|
|
|
func (hio headerIndexObject) Len() int {
|
|
return len(hio)
|
|
}
|
|
|
|
func (hio headerIndexObject) Value(k string) ucl.Object {
|
|
v, ok := hio[k]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
return ucl.IntObject(v)
|
|
}
|
|
|
|
func (hio headerIndexObject) Each(fn func(k string, v ucl.Object) error) error {
|
|
for k, v := range hio {
|
|
if err := fn(k, ucl.IntObject(v)); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|