diff --git a/csv.go b/csv.go index c49469d..3d715f4 100644 --- a/csv.go +++ b/csv.go @@ -3,8 +3,10 @@ package scriptx import ( "bufio" "encoding/csv" + "encoding/json" "errors" "io" + "sort" ) type Splitter interface { @@ -79,3 +81,162 @@ func CSVColumn(name string) func(r io.Reader, w io.Writer) error { } } } + +type CSVHeader struct { + header []string +} + +func (ch *CSVHeader) Column(name string) int { + for c, h := range ch.header { + if h == name { + return c + } + } + return -1 +} + +func (ch *CSVHeader) Value(row []string, name string) string { + c := ch.Column(name) + if c == -1 || c >= len(row) { + return "" + } + return row[c] +} + +func CSVFilter(fn func(row []string, header *CSVHeader) []string) func(io.Reader, io.Writer) error { + return func(r io.Reader, w io.Writer) error { + csvr := csv.NewReader(r) + csvr.FieldsPerRecord = 0 + + csvWriter := csv.NewWriter(w) + + header, rErr := csvr.Read() + if rErr != nil { + return rErr + } + if err := csvWriter.Write(header); err != nil { + return err + } + headerInfo := CSVHeader{header: header} + + row, rErr := csvr.Read() + for !errors.Is(rErr, io.EOF) { + newRow := fn(row, &headerInfo) + if newRow != nil { + if err := csvWriter.Write(newRow); err != nil { + return err + } + } + + row, rErr = csvr.Read() + } + if rErr != nil && !errors.Is(rErr, io.EOF) { + return rErr + } + + csvWriter.Flush() + return nil + } +} + +func CSVSort(lessThan func(row1, row2 []string, header *CSVHeader) bool) func(io.Reader, io.Writer) error { + return func(r io.Reader, w io.Writer) error { + csvr := csv.NewReader(r) + csvr.FieldsPerRecord = 0 + + header, rErr := csvr.Read() + if rErr != nil { + return rErr + } + headerInfo := CSVHeader{header: header} + + records, rErr := csvr.ReadAll() + if rErr != nil { + return rErr + } + + sort.Slice(records, func(i, j int) bool { + return lessThan(records[i], records[j], &headerInfo) + }) + + csvWriter := csv.NewWriter(w) + + if err := csvWriter.Write(header); err != nil { + return err + } + for _, r := range records { + if err := csvWriter.Write(r); err != nil { + return err + } + } + + csvWriter.Flush() + return nil + } +} + +func CSVMapToString(fn func(row []string, header *CSVHeader) string) func(io.Reader, io.Writer) error { + return func(r io.Reader, w io.Writer) error { + csvr := csv.NewReader(r) + csvr.FieldsPerRecord = 0 + + header, rErr := csvr.Read() + if rErr != nil { + return rErr + } + headerInfo := CSVHeader{header: header} + + row, rErr := csvr.Read() + for !errors.Is(rErr, io.EOF) { + line := fn(row, &headerInfo) + if _, err := io.WriteString(w, line); err != nil { + return err + } + if _, err := io.WriteString(w, "\n"); err != nil { + return err + } + + row, rErr = csvr.Read() + } + if rErr != nil && !errors.Is(rErr, io.EOF) { + return rErr + } + + return nil + } +} + +func CSVMapToJSON(fn func(row []string, header *CSVHeader) any) func(io.Reader, io.Writer) error { + return func(r io.Reader, w io.Writer) error { + csvr := csv.NewReader(r) + csvr.FieldsPerRecord = 0 + + header, rErr := csvr.Read() + if rErr != nil { + return rErr + } + headerInfo := CSVHeader{header: header} + + row, rErr := csvr.Read() + for !errors.Is(rErr, io.EOF) { + obj := fn(row, &headerInfo) + jsonBytes, err := json.Marshal(obj) + if err != nil { + return err + } + if _, err = w.Write(jsonBytes); err != nil { + return err + } + if _, err := io.WriteString(w, "\n"); err != nil { + return err + } + + row, rErr = csvr.Read() + } + if rErr != nil && !errors.Is(rErr, io.EOF) { + return rErr + } + + return nil + } +} diff --git a/csv_test.go b/csv_test.go index 93b7175..40cab2e 100644 --- a/csv_test.go +++ b/csv_test.go @@ -57,3 +57,26 @@ func ExampleToCSV() { // b,banana,bravo // c,cherry,charlie } + +func ExampleCSVMapToJSON() { + script.Slice([]string{ + "letter,fruit,word", + "a,apple,alpha", + "b,banana,bravo", + "c,cherry,charlie", + }).Filter(scriptx.CSVMapToJSON(func(row []string, header *scriptx.CSVHeader) any { + return struct { + Letter string `json:"letter"` + Fruit string `json:"fruit"` + Word string `json:"word"` + }{ + Letter: header.Value(row, "letter"), + Fruit: header.Value(row, "fruit"), + Word: header.Value(row, "word"), + } + })).Stdout() + // Output: + // {"letter":"a","fruit":"apple","word":"alpha"} + // {"letter":"b","fruit":"banana","word":"bravo"} + // {"letter":"c","fruit":"cherry","word":"charlie"} +}