This close qax-os#2075, add new function SetColStyle for streaming writer to support set columns style (qax-os#2076)

- Add new exported error variable ErrStreamSetColStyle
- Fix cell default style doesn't override by none-zero row style when set row by stream writer
- Update unit tests

Signed-off-by: mengzhongyuan <mengzhongyuan@bytedance.com>

Author: mengpromax

col.go

@@ -450,6 +450,21 @@ func (f *File) SetColStyle(sheet, columns string, styleID int) error {
 	}
 	s.mu.Unlock()
 	ws.mu.Lock()
+	ws.setColStyle(minVal, maxVal, styleID)
+	ws.mu.Unlock()
+	if rows := len(ws.SheetData.Row); rows > 0 {
+		for col := minVal; col <= maxVal; col++ {
+			from, _ := CoordinatesToCellName(col, 1)
+			to, _ := CoordinatesToCellName(col, rows)
+			err = f.SetCellStyle(sheet, from, to, styleID)
+		}
+	}
+	return err
+}
+
+// setColStyle provides a function to set the style of a single column or
+// multiple columns.
+func (ws *xlsxWorksheet) setColStyle(minVal, maxVal, styleID int) {
 	if ws.Cols == nil {
 		ws.Cols = &xlsxCols{}
 	}
@@ -472,15 +487,6 @@ func (f *File) SetColStyle(sheet, columns string, styleID int) error {
 		fc.Width = c.Width
 		return fc
 	})
-	ws.mu.Unlock()
-	if rows := len(ws.SheetData.Row); rows > 0 {
-		for col := minVal; col <= maxVal; col++ {
-			from, _ := CoordinatesToCellName(col, 1)
-			to, _ := CoordinatesToCellName(col, rows)
-			err = f.SetCellStyle(sheet, from, to, styleID)
-		}
-	}
-	return err
 }
 
 // SetColWidth provides a function to set the width of a single column or
@@ -504,6 +510,13 @@ func (f *File) SetColWidth(sheet, startCol, endCol string, width float64) error
 	f.mu.Unlock()
 	ws.mu.Lock()
 	defer ws.mu.Unlock()
+	ws.setColWidth(minVal, maxVal, width)
+	return err
+}
+
+// setColWidth provides a function to set the width of a single column or
+// multiple columns.
+func (ws *xlsxWorksheet) setColWidth(minVal, maxVal int, width float64) {
 	col := xlsxCol{
 		Min:         minVal,
 		Max:         maxVal,
@@ -514,7 +527,7 @@ func (f *File) SetColWidth(sheet, startCol, endCol string, width float64) error
 		cols := xlsxCols{}
 		cols.Col = append(cols.Col, col)
 		ws.Cols = &cols
-		return err
+		return
 	}
 	ws.Cols.Col = flatCols(col, ws.Cols.Col, func(fc, c xlsxCol) xlsxCol {
 		fc.BestFit = c.BestFit
@@ -525,7 +538,6 @@ func (f *File) SetColWidth(sheet, startCol, endCol string, width float64) error
 		fc.Style = c.Style
 		return fc
 	})
-	return err
 }
 
 // flatCols provides a method for the column's operation functions to flatten

errors.go

@@ -135,6 +135,9 @@ var (
 	// ErrSparklineType defined the error message on receive the invalid
 	// sparkline Type parameters.
 	ErrSparklineType = errors.New("parameter 'Type' must be 'line', 'column' or 'win_loss'")
+	// ErrStreamSetColStyle defined the error message on set column style in
+	// stream writing mode.
+	ErrStreamSetColStyle = errors.New("must call the SetColStyle function before the SetRow function")
 	// ErrStreamSetColWidth defined the error message on set column width in
 	// stream writing mode.
 	ErrStreamSetColWidth = errors.New("must call the SetColWidth function before the SetRow function")

stream.go

@@ -29,7 +29,6 @@ type StreamWriter struct {
 	Sheet           string
 	SheetID         int
 	sheetWritten    bool
-	cols            strings.Builder
 	worksheet       *xlsxWorksheet
 	rawData         bufferedWriter
 	rows            int
@@ -413,16 +412,18 @@ func (sw *StreamWriter) SetRow(cell string, values []interface{}, opts ...RowOpt
 		if err != nil {
 			return err
 		}
-		c := xlsxC{R: ref, S: options.StyleID}
+		c := xlsxC{R: ref, S: sw.worksheet.prepareCellStyle(col, row, options.StyleID)}
+		var s int
 		if v, ok := val.(Cell); ok {
-			c.S = v.StyleID
-			val = v.Value
+			s, val = v.StyleID, v.Value
 			setCellFormula(&c, v.Formula)
 		} else if v, ok := val.(*Cell); ok && v != nil {
-			c.S = v.StyleID
-			val = v.Value
+			s, val = v.StyleID, v.Value
 			setCellFormula(&c, v.Formula)
 		}
+		if s > 0 {
+			c.S = s
+		}
 		if err = sw.setCellValFunc(&c, val); err != nil {
 			_, _ = sw.rawData.WriteString(`</row>`)
 			return err
@@ -433,6 +434,33 @@ func (sw *StreamWriter) SetRow(cell string, values []interface{}, opts ...RowOpt
 	return sw.rawData.Sync()
 }
 
+// SetColStyle provides a function to set the style of a single column or
+// multiple columns for the StreamWriter. Note that you must call
+// the 'SetColStyle' function before the 'SetRow' function. For example set
+// style of column H on Sheet1:
+//
+//	err := sw.SetColStyle(8, 8, style)
+func (sw *StreamWriter) SetColStyle(minVal, maxVal, styleID int) error {
+	if sw.sheetWritten {
+		return ErrStreamSetColStyle
+	}
+	if minVal < MinColumns || minVal > MaxColumns || maxVal < MinColumns || maxVal > MaxColumns {
+		return ErrColumnNumber
+	}
+	if maxVal < minVal {
+		minVal, maxVal = maxVal, minVal
+	}
+	s, err := sw.file.stylesReader()
+	if err != nil {
+		return err
+	}
+	if styleID < 0 || s.CellXfs == nil || len(s.CellXfs.Xf) <= styleID {
+		return newInvalidStyleID(styleID)
+	}
+	sw.worksheet.setColStyle(minVal, maxVal, styleID)
+	return nil
+}
+
 // SetColWidth provides a function to set the width of a single column or
 // multiple columns for the StreamWriter. Note that you must call
 // the 'SetColWidth' function before the 'SetRow' function. For example set
@@ -452,14 +480,7 @@ func (sw *StreamWriter) SetColWidth(minVal, maxVal int, width float64) error {
 	if minVal > maxVal {
 		minVal, maxVal = maxVal, minVal
 	}
-
-	sw.cols.WriteString(`<col min="`)
-	sw.cols.WriteString(strconv.Itoa(minVal))
-	sw.cols.WriteString(`" max="`)
-	sw.cols.WriteString(strconv.Itoa(maxVal))
-	sw.cols.WriteString(`" width="`)
-	sw.cols.WriteString(strconv.FormatFloat(width, 'f', -1, 64))
-	sw.cols.WriteString(`" customWidth="1"/>`)
+	sw.worksheet.setColWidth(minVal, maxVal, width)
 	return nil
 }
 
@@ -642,10 +663,27 @@ func writeCell(buf *bufferedWriter, c xlsxC) {
 func (sw *StreamWriter) writeSheetData() {
 	if !sw.sheetWritten {
 		bulkAppendFields(&sw.rawData, sw.worksheet, 4, 5)
-		if sw.cols.Len() > 0 {
-			_, _ = sw.rawData.WriteString("<cols>")
-			_, _ = sw.rawData.WriteString(sw.cols.String())
-			_, _ = sw.rawData.WriteString("</cols>")
+		if sw.worksheet.Cols != nil {
+			for _, col := range sw.worksheet.Cols.Col {
+				_, _ = sw.rawData.WriteString("<cols>")
+				sw.rawData.WriteString(`<col min="`)
+				sw.rawData.WriteString(strconv.Itoa(col.Min))
+				sw.rawData.WriteString(`" max="`)
+				sw.rawData.WriteString(strconv.Itoa(col.Max))
+				sw.rawData.WriteString(`"`)
+				if col.Width != nil {
+					sw.rawData.WriteString(` width="`)
+					sw.rawData.WriteString(strconv.FormatFloat(*col.Width, 'f', -1, 64))
+					sw.rawData.WriteString(`" customWidth="1"`)
+				}
+				if col.Style != 0 {
+					sw.rawData.WriteString(` style="`)
+					sw.rawData.WriteString(strconv.Itoa(col.Style))
+					sw.rawData.WriteString(`"`)
+				}
+				sw.rawData.WriteString(`/>`)
+				_, _ = sw.rawData.WriteString("</cols>")
+			}
 		}
 		_, _ = sw.rawData.WriteString(`<sheetData>`)
 		sw.sheetWritten = true

stream_test.go

@@ -154,19 +154,53 @@ func TestStreamWriter(t *testing.T) {
 	assert.NoError(t, file.Close())
 }
 
+func TestStreamSetColStyle(t *testing.T) {
+	file := NewFile()
+	defer func() {
+		assert.NoError(t, file.Close())
+	}()
+	streamWriter, err := file.NewStreamWriter("Sheet1")
+	assert.NoError(t, err)
+	assert.NoError(t, streamWriter.SetColStyle(3, 2, 0))
+	assert.Equal(t, ErrColumnNumber, streamWriter.SetColStyle(0, 3, 20))
+	assert.Equal(t, ErrColumnNumber, streamWriter.SetColStyle(MaxColumns+1, 3, 20))
+	assert.Equal(t, newInvalidStyleID(2), streamWriter.SetColStyle(1, 3, 2))
+	assert.NoError(t, streamWriter.SetRow("A1", []interface{}{"A", "B", "C"}))
+	assert.Equal(t, ErrStreamSetColStyle, streamWriter.SetColStyle(2, 3, 0))
+
+	file = NewFile()
+	defer func() {
+		assert.NoError(t, file.Close())
+	}()
+	// Test set column style with unsupported charset style sheet
+	file.Styles = nil
+	file.Pkg.Store(defaultXMLPathStyles, MacintoshCyrillicCharset)
+	streamWriter, err = file.NewStreamWriter("Sheet1")
+	assert.NoError(t, err)
+	assert.EqualError(t, streamWriter.SetColStyle(3, 2, 0), "XML syntax error on line 1: invalid UTF-8")
+}
+
 func TestStreamSetColWidth(t *testing.T) {
 	file := NewFile()
 	defer func() {
 		assert.NoError(t, file.Close())
 	}()
+	styleID, err := file.NewStyle(&Style{
+		Fill: Fill{Type: "pattern", Color: []string{"E0EBF5"}, Pattern: 1},
+	})
+	if err != nil {
+		fmt.Println(err)
+	}
 	streamWriter, err := file.NewStreamWriter("Sheet1")
 	assert.NoError(t, err)
 	assert.NoError(t, streamWriter.SetColWidth(3, 2, 20))
+	assert.NoError(t, streamWriter.SetColStyle(3, 2, styleID))
 	assert.Equal(t, ErrColumnNumber, streamWriter.SetColWidth(0, 3, 20))
 	assert.Equal(t, ErrColumnNumber, streamWriter.SetColWidth(MaxColumns+1, 3, 20))
 	assert.Equal(t, ErrColumnWidth, streamWriter.SetColWidth(1, 3, MaxColumnWidth+1))
 	assert.NoError(t, streamWriter.SetRow("A1", []interface{}{"A", "B", "C"}))
 	assert.Equal(t, ErrStreamSetColWidth, streamWriter.SetColWidth(2, 3, 20))
+	assert.NoError(t, streamWriter.Flush())
 }
 
 func TestStreamSetPanes(t *testing.T) {
@@ -323,7 +357,6 @@ func TestStreamSetRowWithStyle(t *testing.T) {
 	defer func() {
 		assert.NoError(t, file.Close())
 	}()
-	zeroStyleID := 0
 	grayStyleID, err := file.NewStyle(&Style{Font: &Font{Color: "777777"}})
 	assert.NoError(t, err)
 	blueStyleID, err := file.NewStyle(&Style{Font: &Font{Color: "0000FF"}})
@@ -342,7 +375,7 @@ func TestStreamSetRowWithStyle(t *testing.T) {
 
 	ws, err := file.workSheetReader("Sheet1")
 	assert.NoError(t, err)
-	for colIdx, expected := range []int{grayStyleID, zeroStyleID, zeroStyleID, blueStyleID, blueStyleID} {
+	for colIdx, expected := range []int{grayStyleID, grayStyleID, grayStyleID, blueStyleID, blueStyleID} {
 		assert.Equal(t, expected, ws.SheetData.Row[0].C[colIdx].S)
 	}
 }
@@ -381,8 +414,8 @@ func TestStreamSetCellValFunc(t *testing.T) {
 }
 
 func TestSetCellIntFunc(t *testing.T) {
-	cases := []struct{
-		val interface{}
+	cases := []struct {
+		val    interface{}
 		target string
 	}{
 		{val: 128, target: "128"},