From c52dc2b3e986424e6b86ba7743255fe115eaddf4 Mon Sep 17 00:00:00 2001 From: Leon Mika Date: Mon, 3 Feb 2025 10:28:44 +1100 Subject: [PATCH] Added csv:to-csv --- _docs/mod/csv.md | 13 +++- ucl/builtins.go | 16 ++--- ucl/builtins/csv.go | 126 +++++++++++++++++++++++++++++++++++++++ ucl/builtins/csv_test.go | 45 ++++++++++++++ ucl/objs.go | 2 +- ucl/userbuiltin.go | 6 ++ 6 files changed, 198 insertions(+), 10 deletions(-) diff --git a/_docs/mod/csv.md b/_docs/mod/csv.md index e630c0c..8926ad1 100644 --- a/_docs/mod/csv.md +++ b/_docs/mod/csv.md @@ -26,4 +26,15 @@ csv:each-record "winds.csv" { |row hdr| echo "Wind $name has bearing $bearing" } -``` \ No newline at end of file +``` + +``` +csv:to-csv CONT [-header HEADER] +``` + +Produces a CSV using the items of CONT and writes it as a string. CONT must be a list of iterator of hashable +elements. + +If HEADER is defined, it must be a list of strings identifying the name and order of the hash keys to read +from each hashable item of CONT. If HEADER is not defined, then the keys of the first consumed hashable item will be +used, with the keys sorted in alphabetical order. \ No newline at end of file diff --git a/ucl/builtins.go b/ucl/builtins.go index b81eaf6..332a546 100644 --- a/ucl/builtins.go +++ b/ucl/builtins.go @@ -340,8 +340,8 @@ func objectsEqual(l, r Object) bool { } } return true - case hashable: - rv, ok := r.(hashable) + case Hashable: + rv, ok := r.(Hashable) if !ok { return false } @@ -462,7 +462,7 @@ func lenBuiltin(ctx context.Context, args invocationArgs) (Object, error) { return IntObject(len(string(v))), nil case Listable: return IntObject(v.Len()), nil - case hashable: + case Hashable: return IntObject(v.Len()), nil case Iterable: cnt := 0 @@ -490,7 +490,7 @@ func indexLookup(ctx context.Context, obj, elem Object) (Object, error) { return v.Index(int(intIdx)), nil } return nil, nil - case hashable: + case Hashable: strIdx, ok := elem.(StringObject) if !ok { return nil, errors.New("expected string for hashable") @@ -524,7 +524,7 @@ func keysBuiltin(ctx context.Context, args invocationArgs) (Object, error) { val := args.args[0] switch v := val.(type) { - case hashable: + case Hashable: keys := make(ListObject, 0, v.Len()) if err := v.Each(func(k string, _ Object) error { keys = append(keys, StringObject(k)) @@ -669,7 +669,7 @@ func filterBuiltin(ctx context.Context, args invocationArgs) (Object, error) { } } return &newList, nil - case hashable: + case Hashable: newHash := hashObject{} if err := t.Each(func(k string, v Object) error { if m, err := inv.invoke(ctx, args.fork([]Object{StringObject(k), v})); err != nil { @@ -734,7 +734,7 @@ func reduceBuiltin(ctx context.Context, args invocationArgs) (Object, error) { accum = newAccum } return accum, nil - case hashable: + case Hashable: // TODO: should raise error? if err := t.Each(func(k string, v Object) error { newAccum, err := block.invoke(ctx, args.fork([]Object{StringObject(k), v, accum})) @@ -1031,7 +1031,7 @@ func foreachBuiltin(ctx context.Context, args macroArgs) (Object, error) { } } } - case hashable: + case Hashable: err := t.Each(func(k string, v Object) error { last, err = args.evalBlock(ctx, blockIdx, []Object{StringObject(k), v}, true) return err diff --git a/ucl/builtins/csv.go b/ucl/builtins/csv.go index 95b30d5..c013549 100644 --- a/ucl/builtins/csv.go +++ b/ucl/builtins/csv.go @@ -4,9 +4,11 @@ import ( "context" "encoding/csv" "errors" + "fmt" "io" "io/fs" "os" + "sort" "strings" "ucl.lmika.dev/ucl" ) @@ -22,6 +24,7 @@ func CSV(fs fs.FS) ucl.Module { Name: "csv", Builtins: map[string]ucl.BuiltinHandler{ "each-record": fsh.eachRecord, + "to-csv": fsh.toCsv, }, } } @@ -76,6 +79,129 @@ func (h csvHandlers) eachRecord(ctx context.Context, args ucl.CallArgs) (any, er return nil, nil } +func (h csvHandlers) toCsv(ctx context.Context, args ucl.CallArgs) (any, error) { + var rowConsumer func() (ucl.Object, error) + + var ( + listable ucl.Listable + iterable ucl.Iterable + ) + if args.CanBind(&listable) { + if err := args.Bind(&listable); err != nil { + return nil, err + } + + rowConsumer = func() func() (ucl.Object, error) { + idx := 0 + return func() (ucl.Object, error) { + if idx >= listable.Len() { + return nil, io.EOF + } + + v := listable.Index(idx) + idx += 1 + return v, nil + } + }() + } else if args.CanBind(&iterable) { + if err := args.Bind(&iterable); err != nil { + return nil, err + } + + rowConsumer = func() func() (ucl.Object, error) { + return func() (ucl.Object, error) { + if !iterable.HasNext() { + return nil, io.EOF + } + + return iterable.Next(ctx) + } + }() + } else { + return nil, errors.New("unsupported type") + } + + var ( + header []string + row []string + ) + if args.HasSwitch("header") { + var headerListable ucl.Listable + + if err := args.BindSwitch("header", &headerListable); err != nil { + return nil, err + } + if headerListable.Len() == 0 { + return nil, errors.New("header requires at least one item") + } + + header = make([]string, headerListable.Len()) + for i := range header { + header[i] = headerListable.Index(i).String() + } + + row = make([]string, len(header)) + } + + var sb strings.Builder + cw := csv.NewWriter(&sb) + + i := 0 + for { + v, err := rowConsumer() + if err != nil { + if errors.Is(err, io.EOF) { + break + } else { + return nil, err + } + } + + hv, ok := v.(ucl.Hashable) + if !ok { + return nil, fmt.Errorf("element %v is not hashable", i) + } + + if len(header) == 0 { + if err := hv.Each(func(k string, _ ucl.Object) error { + header = append(header, k) + return nil + }); err != nil { + return nil, err + } + + sort.Strings(header) + row = make([]string, len(header)) + } + if i == 0 { + if err := cw.Write(header); err != nil { + return nil, err + } + } + + for j, k := range header { + val := hv.Value(k) + if val == nil { + row[j] = "" + } else { + row[j] = val.String() + } + } + + if err := cw.Write(row); err != nil { + return nil, err + } + + i++ + } + + cw.Flush() + if err := cw.Error(); err != nil { + return nil, err + } + return sb.String(), nil +} + type headerIndexObject map[string]int func (hio headerIndexObject) String() string { diff --git a/ucl/builtins/csv_test.go b/ucl/builtins/csv_test.go index 5b60834..e8a436e 100644 --- a/ucl/builtins/csv_test.go +++ b/ucl/builtins/csv_test.go @@ -51,3 +51,48 @@ func TestCSV_ReadRecord(t *testing.T) { }) } } + +func TestCSV_ToCSV(t *testing.T) { + tests := []struct { + descr string + eval string + want string + }{ + { + descr: "to csv 1", + eval: `csv:to-csv [[hello:"one"] [hello:"two"]]`, + want: "hello\none\ntwo\n", + }, + { + descr: "to csv 2", + eval: `csv:to-csv [[hello:"one" world:"this"] [hello:"two" world:"that"]]`, + want: "hello,world\none,this\ntwo,that\n", + }, + { + descr: "to csv 3", + eval: `csv:to-csv [[hello:"one" world:"this"] [hello:"two" other:"the" world:"that"]] -header [world other]`, + want: "world,other\nthis,\nthat,the\n", + }, + { + descr: "to csv 4", + eval: `itrs:from [[hello:"one" world:"this"] [hello:"two" other:"the" world:"that"]] | csv:to-csv -header [world other]`, + want: "world,other\nthis,\nthat,the\n", + }, + } + + for _, tt := range tests { + t.Run(tt.descr, func(t *testing.T) { + var bfr bytes.Buffer + + inst := ucl.New( + ucl.WithModule(builtins.CSV(testCsvFS)), + ucl.WithModule(builtins.Itrs()), + ucl.WithOut(&bfr), + ) + + res, err := inst.Eval(context.Background(), tt.eval) + assert.NoError(t, err) + assert.Equal(t, tt.want, res) + }) + } +} diff --git a/ucl/objs.go b/ucl/objs.go index 149c965..bf0cca8 100644 --- a/ucl/objs.go +++ b/ucl/objs.go @@ -40,7 +40,7 @@ type ModListable interface { Insert(idx int, obj Object) error } -type hashable interface { +type Hashable interface { Len() int Value(k string) Object Each(func(k string, v Object) error) error diff --git a/ucl/userbuiltin.go b/ucl/userbuiltin.go index 2d74a57..f09fe45 100644 --- a/ucl/userbuiltin.go +++ b/ucl/userbuiltin.go @@ -186,6 +186,12 @@ func canBindArg(v interface{}, arg Object) bool { case *int: _, ok := arg.(IntObject) return ok + case *Listable: + _, ok := arg.(Listable) + return ok + case *Iterable: + _, ok := arg.(Iterable) + return ok } switch t := arg.(type) {