HeaderModifier (#10)

* feat: HeaderModifier * feat: use HeaderModifier * feat: add test for HeaderModifier * fix: HeaderModifier in XsvWrite * refactor: delete for loop nest
shigetaichi · Sep 4, 2023 · 46922c1 · 46922c1
1 parent 0561755
commit 46922c1
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 4 deletions.
diff --git a/encode_test.go b/encode_test.go
@@ -708,3 +708,32 @@ func Test_writeToChan_emptyptr_sortOrder(t *testing.T) {
 		assertLine(t, []string{strconv.Itoa(i - 1), "f", "baz" + strconv.Itoa(i-1), strconv.FormatFloat(float64(i-1), 'f', -1, 64), "", "*string", ""}, l)
 	}
 }
+
+func Test_HeaderModifier(t *testing.T) {
+	b := bytes.Buffer{}
+	e := &encoder{out: &b}
+	blah := 2
+	sptr := "*string"
+	s := []Sample{
+		{Foo: "f", Bar: 1, Baz: "baz", Frop: 0.1, Blah: &blah, SPtr: &sptr},
+		{Foo: "e", Bar: 3, Baz: "b", Frop: 6.0 / 13, Blah: nil, SPtr: nil},
+	}
+
+	xsvWrite := NewXsvWrite[Sample]()
+	xsvWrite.HeaderModifier = map[string]string{"BAR": "BAR-updated"}
+	xsvWrite.HeaderModifier["Omit"] = ""
+	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) != 3 {
+		t.Fatalf("expected 3 lines, got %d", len(lines))
+	}
+	assertLine(t, []string{"foo", "BAR-updated", "Baz", "Quux", "Blah", "SPtr", ""}, lines[0])
+	assertLine(t, []string{"f", "1", "baz", "0.1", "2", "*string", ""}, lines[1])
+	assertLine(t, []string{"e", "3", "b", "0.46153846153846156", "", "", ""}, lines[2])
+}
diff --git a/xsv_write.go b/xsv_write.go
@@ -13,8 +13,9 @@ 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
-	SortOrder       []int    // column sort order
+	SelectedColumns []string          // slice of field names to output
+	SortOrder       []int             // column sort order
+	HeaderModifier  map[string]string // map to dynamically change headers
 	nameNormalizer  Normalizer
 }
 
@@ -25,6 +26,7 @@ func NewXsvWrite[T any]() XsvWrite[T] {
 		OmitHeaders:     false,
 		SelectedColumns: make([]string, 0),
 		SortOrder:       make([]int, 0),
+		HeaderModifier:  map[string]string{},
 		nameNormalizer:  func(s string) string { return s },
 	}
 }

diff --git a/xsv_writer.go b/xsv_writer.go
@@ -42,7 +42,11 @@ func (xw *XsvWriter[T]) Write(data []T) error {
 
 	csvHeadersLabels := make([]string, len(inInnerStructInfo.Fields))
 	for i, fieldInfo := range inInnerStructInfo.Fields { // Used to write the header (first line) in CSV
-		csvHeadersLabels[i] = fieldInfo.getFirstKey()
+		if newHeader, ok := xw.HeaderModifier[fieldInfo.getFirstKey()]; ok { // modify header name dynamically
+			csvHeadersLabels[i] = newHeader
+		} else {
+			csvHeadersLabels[i] = fieldInfo.getFirstKey()
+		}
 	}
 	if !xw.OmitHeaders {
 		if err := xw.writer.Write(csvHeadersLabels); err != nil {
@@ -87,7 +91,11 @@ func (xw *XsvWriter[T]) WriteFromChan(dataChan chan T) error {
 	inInnerStructInfo := &structInfo{fieldInfos}
 	csvHeadersLabels := make([]string, len(inInnerStructInfo.Fields))
 	for i, fieldInfo := range inInnerStructInfo.Fields { // Used to Write the header (first line) in CSV
-		csvHeadersLabels[i] = fieldInfo.getFirstKey()
+		if newHeader, ok := xw.HeaderModifier[fieldInfo.getFirstKey()]; ok { // modify header name dynamically
+			csvHeadersLabels[i] = newHeader
+		} else {
+			csvHeadersLabels[i] = fieldInfo.getFirstKey()
+		}
 	}
 
 	if !xw.OmitHeaders {