SortOrder (#7)

* feat: add Test_writeTo_emptyptr_sortOrder * feat: add SortOrder * fix: checkSortOrderSlice
shigetaichi · Aug 15, 2023 · a29b8ed · a29b8ed
1 parent c4354d6
commit a29b8ed
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 10 deletions.
diff --git a/encode_test.go b/encode_test.go
@@ -640,3 +640,36 @@ func Test_writeTo_emptyptr_selectedColumns(t *testing.T) {
 	assertLine(t, []string{"first", "Baz", "last"}, lines[0])
 	assertLine(t, []string{"aaa", "baz", "zzz"}, lines[1])
 }
+
+func Test_writeTo_emptyptr_sortOrder(t *testing.T) {
+
+	b := bytes.Buffer{}
+	e := &encoder{out: &b}
+	blah := 2
+	sptr := "*string"
+	s := []EmbedPtrSample{
+		{
+			Qux:    "aaa",
+			Sample: &Sample{Foo: "f", Bar: 1, Baz: "baz", Frop: 0.2, Blah: &blah, SPtr: &sptr},
+			Ignore: "shouldn't be marshalled",
+			Quux:   "zzz",
+			Grault: math.Pi,
+		},
+	}
+
+	xsvWrite := NewXsvWrite[EmbedPtrSample]()
+	xsvWrite.SortOrder = []int{1, 0, 2, 3, 4, 5, 6, 7, 8, 9}
+	if err := xsvWrite.SetWriter(csv.NewWriter(e.out)).Write(s); err != nil {
+		t.Fatal(err)
+	}
+
+	lines, err := csv.NewReader(&b).ReadAll()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(lines) != 2 {
+		t.Fatalf("expected 2 lines, got %d", len(lines))
+	}
+	assertLine(t, []string{"foo", "first", "BAR", "Baz", "Quux", "Blah", "SPtr", "Omit", "garply", "last"}, lines[0])
+	assertLine(t, []string{"f", "aaa", "1", "baz", "0.2", "2", "*string", "", "3.141592653589793", "zzz"}, lines[1])
+}
diff --git a/xsv_write.go b/xsv_write.go
@@ -3,6 +3,8 @@ package xsv
 import (
 	"bytes"
 	"encoding/csv"
+	"errors"
+	"fmt"
 	"os"
 	"slices"
 )
@@ -11,25 +13,31 @@ type XsvWrite[T any] struct {
 	TagName         string //key in the struct field's tag to scan
 	TagSeparator    string //separator string for multiple csv tags in struct fields
 	OmitHeaders     bool
-	SelectedColumns []string     // slice of field names to output
-	columnSorter    ColumnSorter // TODO: describe in comment
+	SelectedColumns []string // slice of field names to output
+	SortOrder       []int    // column sort order
 	nameNormalizer  Normalizer
 }
-type ColumnSorter = func(row []string) []string
 
 func NewXsvWrite[T any]() XsvWrite[T] {
 	return XsvWrite[T]{
 		TagName:         "csv",
 		TagSeparator:    ",",
 		OmitHeaders:     false,
 		SelectedColumns: make([]string, 0),
-		columnSorter: func(row []string) []string {
-			return row
-		},
-		nameNormalizer: func(s string) string { return s },
+		SortOrder:       make([]int, 0),
+		nameNormalizer:  func(s string) string { return s },
 	}
 }
 
+func (x *XsvWrite[T]) checkSortOrderSlice(outputFieldsCount int) error {
+	if len(x.SortOrder) > 0 {
+		if len(x.SortOrder) != outputFieldsCount {
+			return errors.New(fmt.Sprintf("the length of the SortOrder array should be equal to the number of items to be output(%d)", outputFieldsCount))
+		}
+	}
+	return nil
+}
+
 func (x *XsvWrite[T]) getSelectedFieldInfos(fieldInfos []fieldInfo) []fieldInfo {
 	if len(x.SelectedColumns) > 0 {
 		var selectedFieldInfos []fieldInfo

diff --git a/xsv_writer.go b/xsv_writer.go
@@ -40,7 +40,10 @@ func (xw *XsvWriter[T]) Write(data []T) error {
 	for i, fieldInfo := range inInnerStructInfo.Fields { // Used to write the header (first line) in CSV
 		csvHeadersLabels[i] = fieldInfo.getFirstKey()
 	}
-	csvHeadersLabels = xw.columnSorter(csvHeadersLabels)
+	if err := xw.checkSortOrderSlice(len(fieldInfos)); err != nil {
+		return err
+	}
+	csvHeadersLabels = reorderColumns(csvHeadersLabels, xw.SortOrder)
 	if !xw.OmitHeaders {
 		if err := xw.writer.Write(csvHeadersLabels); err != nil {
 			return err
@@ -56,7 +59,7 @@ func (xw *XsvWriter[T]) Write(data []T) error {
 			}
 			csvHeadersLabels[j] = inInnerFieldValue
 		}
-		csvHeadersLabels = xw.columnSorter(csvHeadersLabels)
+		csvHeadersLabels = reorderColumns(csvHeadersLabels, xw.SortOrder)
 		if err := xw.writer.Write(csvHeadersLabels); err != nil {
 			return err
 		}
@@ -83,6 +86,11 @@ func (xw *XsvWriter[T]) WriteFromChan(dataChan chan T) error {
 	for i, fieldInfo := range inInnerStructInfo.Fields { // Used to Write the header (first line) in CSV
 		csvHeadersLabels[i] = fieldInfo.getFirstKey()
 	}
+
+	if err := xw.checkSortOrderSlice(len(fieldInfos)); err != nil {
+		return err
+	}
+	csvHeadersLabels = reorderColumns(csvHeadersLabels, xw.SortOrder)
 	if !xw.OmitHeaders {
 		if err := xw.writer.Write(csvHeadersLabels); err != nil {
 			return err
@@ -96,7 +104,7 @@ func (xw *XsvWriter[T]) WriteFromChan(dataChan chan T) error {
 				return err
 			}
 			csvHeadersLabels[j] = inInnerFieldValue
-			csvHeadersLabels = xw.columnSorter(csvHeadersLabels)
+			csvHeadersLabels = reorderColumns(csvHeadersLabels, xw.SortOrder)
 		}
 		if err := xw.writer.Write(csvHeadersLabels); err != nil {
 			return err
@@ -118,3 +126,15 @@ func (xw *XsvWriter[T]) WriteFromChan(dataChan chan T) error {
 	xw.writer.Flush()
 	return xw.writer.Error()
 }
+
+func reorderColumns(row []string, sortOrder []int) []string {
+	if len(sortOrder) > 0 {
+		newLine := make([]string, len(row))
+		for from, to := range sortOrder {
+			newLine[to] = row[from]
+		}
+		return newLine
+	} else {
+		return row
+	}
+}