From a29b8eda745577b558c03fe2b011007998521b9f Mon Sep 17 00:00:00 2001
From: Taichi Shigematsu <65220900+shigetaichi@users.noreply.github.com>
Date: Wed, 16 Aug 2023 02:02:00 +0900
Subject: [PATCH] SortOrder (#7)

* feat: add Test_writeTo_emptyptr_sortOrder

* feat: add SortOrder

* fix: checkSortOrderSlice
---
 encode_test.go | 33 +++++++++++++++++++++++++++++++++
 xsv_write.go   | 22 +++++++++++++++-------
 xsv_writer.go  | 26 +++++++++++++++++++++++---
 3 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/encode_test.go b/encode_test.go
index b6f0367..b7358d3 100644
--- 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
index 5d8bf12..a7212e3 100644
--- a/xsv_write.go
+++ b/xsv_write.go
@@ -3,6 +3,8 @@ package xsv
 import (
 	"bytes"
 	"encoding/csv"
+	"errors"
+	"fmt"
 	"os"
 	"slices"
 )
@@ -11,11 +13,10 @@ 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]{
@@ -23,13 +24,20 @@ func NewXsvWrite[T any]() XsvWrite[T] {
 		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
index e8842fa..77ad2d6 100644
--- 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
+	}
+}