lib/csv/csv.go
// Package csv reads comma-separated values from strings and writes CSV data to strings.
//
// Migrated from https://github.com/qri-io/starlib/tree/master/encoding/csv
package csv
import (
"bytes"
"encoding/csv"
"fmt"
"strings"
"sync"
"github.com/1set/starlet/dataconv"
tps "github.com/1set/starlet/dataconv/types"
"github.com/1set/starlet/internal/replacecr"
"go.starlark.net/starlark"
"go.starlark.net/starlarkstruct"
)
// ModuleName defines the expected name for this Module when used in starlark's load() function, eg: load('csv', 'read_all')
const ModuleName = "csv"
var (
once sync.Once
csvModule starlark.StringDict
)
// LoadModule loads the base64 module.
// It is concurrency-safe and idempotent.
func LoadModule() (starlark.StringDict, error) {
once.Do(func() {
csvModule = starlark.StringDict{
ModuleName: &starlarkstruct.Module{
Name: ModuleName,
Members: starlark.StringDict{
"read_all": starlark.NewBuiltin(ModuleName+".read_all", readAll),
"write_all": starlark.NewBuiltin(ModuleName+".write_all", writeAll),
"write_dict": starlark.NewBuiltin(ModuleName+".write_dict", writeDict),
},
},
}
})
return csvModule, nil
}
// readAll gets all values from a csv source string.
func readAll(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var (
source tps.StringOrBytes
lazyQuotes, trimLeadingSpace bool
skipRow, limitRow int
fieldsPerRecord int
_comma, _comment starlark.String
)
if err := starlark.UnpackArgs(b.Name(), args, kwargs,
"source", &source,
"comma?", &_comma,
"comment", &_comment,
"lazy_quotes", &lazyQuotes,
"trim_leading_space", &trimLeadingSpace,
"fields_per_record?", &fieldsPerRecord,
"skip?", &skipRow,
"limit?", &limitRow); err != nil {
return nil, err
}
// prepare reader
csvr := csv.NewReader(replacecr.Reader(strings.NewReader(source.GoString())))
csvr.LazyQuotes = lazyQuotes
csvr.TrimLeadingSpace = trimLeadingSpace
comma := string(_comma)
if comma == "" {
comma = ","
} else if len(comma) != 1 {
return starlark.None, fmt.Errorf("%s: expected comma param to be a single-character string", b.Name())
}
csvr.Comma = []rune(comma)[0]
comment := string(_comment)
if comment != "" && len(comment) != 1 {
return starlark.None, fmt.Errorf("%s: expected comment param to be a single-character string", b.Name())
} else if comment != "" {
csvr.Comment = []rune(comment)[0]
}
csvr.FieldsPerRecord = fieldsPerRecord
// pre-read to skip rows
if skipRow > 0 {
for i := 0; i < skipRow; i++ {
if _, err := csvr.Read(); err != nil {
return starlark.None, fmt.Errorf("%s: %w", b.Name(), err)
}
}
}
// read all rows
strs, err := csvr.ReadAll()
if err != nil {
return starlark.None, fmt.Errorf("%s: %w", b.Name(), err)
}
// convert and limit rows
vals := make([]starlark.Value, 0, len(strs))
for i, rowStr := range strs {
if limitRow > 0 && i >= limitRow {
break
}
row := make([]starlark.Value, len(rowStr))
for j, cell := range rowStr {
row[j] = starlark.String(cell)
}
vals = append(vals, starlark.NewList(row))
}
return starlark.NewList(vals), nil
}
// writeAll writes a list of lists to a csv string.
func writeAll(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var (
buf = &bytes.Buffer{}
data starlark.Value
comma string
)
if err := starlark.UnpackArgs(b.Name(), args, kwargs, "data", &data, "comma?", &comma); err != nil {
return nil, err
}
// prepare writer
csvw := csv.NewWriter(buf)
if comma == "" {
comma = ","
} else if len(comma) != 1 {
return starlark.None, fmt.Errorf("%s: expected comma param to be a single-character string", b.Name())
}
csvw.Comma = []rune(comma)[0]
// convert data to [][]string
val, err := dataconv.Unmarshal(data)
if err != nil {
return starlark.None, fmt.Errorf("%s: %w", b.Name(), err)
}
sl, ok := val.([]interface{})
if !ok {
return starlark.None, fmt.Errorf("%s: expected value to be an array type", b.Name())
}
var records [][]string
for i, v := range sl {
sl, ok := v.([]interface{})
if !ok {
return starlark.None, fmt.Errorf("%s: row %d is not an array type", b.Name(), i)
}
var row = make([]string, len(sl))
for j, v := range sl {
row[j] = fmt.Sprintf("%v", v)
}
records = append(records, row)
}
// write all records
if err := csvw.WriteAll(records); err != nil {
return starlark.None, fmt.Errorf("%s: %w", b.Name(), err)
}
return starlark.String(buf.String()), nil
}
// writeDict writes a list of dictionaries to a csv string.
func writeDict(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var (
buf = &bytes.Buffer{}
data starlark.Value
header starlark.Iterable
comma string
)
if err := starlark.UnpackArgs(b.Name(), args, kwargs, "data", &data, "header", &header, "comma?", &comma); err != nil {
return nil, err
}
// prepare writer
csvw := csv.NewWriter(buf)
if comma == "" {
comma = ","
} else if len(comma) != 1 {
return starlark.None, fmt.Errorf("%s: expected comma param to be a single-character string", b.Name())
}
csvw.Comma = []rune(comma)[0]
// convert header to []string
var headerStr []string
iter := header.Iterate()
defer iter.Done()
var hv starlark.Value
for iter.Next(&hv) {
s, ok := starlark.AsString(hv)
if !ok {
return starlark.None, fmt.Errorf("%s: for parameter header: got %s, want string", b.Name(), hv.Type())
}
headerStr = append(headerStr, s)
}
if len(headerStr) == 0 {
return starlark.None, fmt.Errorf("%s: header cannot be empty", b.Name())
}
// convert data to []map[string]interface{}
val, err := dataconv.Unmarshal(data)
if err != nil {
return starlark.None, fmt.Errorf("%s: %w", b.Name(), err)
}
sl, ok := val.([]interface{})
if !ok {
return starlark.None, fmt.Errorf("%s: expected value to be an array type", b.Name())
}
// write header
var records [][]string
records = append(records, headerStr)
for _, m := range sl {
// cast to map
mm, ok := m.(map[string]interface{})
if !ok {
return starlark.None, fmt.Errorf("%s: expected value to be a map type", b.Name())
}
// write row
var row = make([]string, len(headerStr))
for j, k := range headerStr {
if v, ok := mm[k]; ok {
row[j] = fmt.Sprintf("%v", v)
}
}
records = append(records, row)
}
// write all records
if err := csvw.WriteAll(records); err != nil {
return starlark.None, fmt.Errorf("%s: %w", b.Name(), err)
}
return starlark.String(buf.String()), nil
}