Cold_Logistic/internal/pkg/utils/excelutil/writer.go

package excelutil

import (
	"context"
	"os"
	"path/filepath"
	"time"

	"gogs.baozhida.cn/Cold_Logistic_libs/pkg/contrib/errors"
)

type DynamicWriteSheet struct {
	File                *excelize.File
	StreamWriter        *excelize.StreamWriter
	Style               *style
	DefaultTitleStyleId int
	DefaultValueStyleId int
	DefaultTimeStyleId  int
	SaveFilePath        string
	SheetName           string
}

type (
	DynamicWriteSheetFun    func(*DynamicWriteSheet) error
	DynamicPreWriteSheetFun func(*DynamicWriteSheet) error
)

func WithDynamicWriteSheetSaveFilePath(saveFilePath string) DynamicWriteSheetFun {
	return func(writer *DynamicWriteSheet) error {
		writer.SaveFilePath = saveFilePath
		return nil
	}
}

func WithDynamicWriteSheetName(sheetName string) DynamicWriteSheetFun {
	return func(writer *DynamicWriteSheet) error {
		writer.SheetName = sheetName
		return nil
	}
}

func NewDynamicWriteSheet(opts ...DynamicWriteSheetFun) (writer *DynamicWriteSheet, err error) {
	writer = &DynamicWriteSheet{
		File:      excelize.NewFile(),
		SheetName: "Sheet1",
	}
	writer.Style = NewStyle(writer.File)
	writer.DefaultTitleStyleId, err = writer.Style.NewBoldTitleStyleId()
	if err != nil {
		return nil, err
	}
	writer.DefaultValueStyleId, err = writer.Style.NewDefaultValueStyleId()
	if err != nil {
		return nil, err
	}
	writer.DefaultTimeStyleId, err = writer.Style.NewDefaultTimeStyleId()
	if err != nil {
		return nil, err
	}
	for _, opt := range opts {
		if err = opt(writer); err != nil {
			return nil, err
		}
	}
	writer.StreamWriter, err = writer.File.NewStreamWriter(writer.SheetName)
	if err != nil {
		return nil, err
	}
	return writer, nil
}

func (w *DynamicWriteSheet) PreWrite(_ context.Context, opts ...DynamicPreWriteSheetFun) error {
	for _, opt := range opts {
		if err := opt(w); err != nil {
			return err
		}
	}
	return nil
}

// WriteToFile 将标题内容写入到文件
func (w *DynamicWriteSheet) WriteToFile(ctx context.Context, titles, contents [][]interface{}) error {
	if len(w.SaveFilePath) == 0 {
		return errors.New("SaveFilePath is empty")
	}
	err := w.WriteTitle(ctx, 0, titles)
	if err != nil {
		return errors.WithStackOnce(err)
	}
	err = w.WriteContent(ctx, len(titles), contents)
	if err != nil {
		return errors.WithStackOnce(err)
	}
	return w.FlushToFile()
}

// WriteToBuff 将标题内容写入到 *excelize.File 缓存中
func (w *DynamicWriteSheet) WriteToBuff(ctx context.Context, titles, contents [][]interface{}) error {
	err := w.WriteTitle(ctx, 0, titles)
	if err != nil {
		return errors.WithStackOnce(err)
	}
	err = w.WriteContent(ctx, len(titles), contents)
	if err != nil {
		return errors.WithStackOnce(err)
	}
	return w.Flush()
}

func (w *DynamicWriteSheet) WriteTitle(ctx context.Context, startRow int, titles [][]interface{}) error {
	return w.Write(ctx, startRow, titles, w.DefaultTitleStyleId)
}

func (w *DynamicWriteSheet) WriteContent(ctx context.Context, startRow int, contents [][]interface{}) error {
	return w.Write(ctx, startRow, contents, w.DefaultValueStyleId)
}

func (w *DynamicWriteSheet) Write(_ context.Context, startRow int, values [][]interface{}, styleId int) (err error) {
	for rowID := 0; rowID < len(values); rowID++ {
		row := values[rowID]
		for i, v := range row {
			cell := excelize.Cell{}
			if val, ok := v.(excelize.Cell); ok {
				cell = val
			} else if _, ok := v.(time.Time); ok {
				cell = excelize.Cell{StyleID: w.DefaultTimeStyleId, Value: v}
			} else {
				cell = excelize.Cell{StyleID: styleId, Value: v}
			}
			row[i] = cell
		}
		cell, _ := excelize.CoordinatesToCellName(1, rowID+startRow+1)
		if err := w.StreamWriter.SetRow(cell, row); err != nil {
			return errors.WithStack(err)
		}
	}
	return nil
}

func (w *DynamicWriteSheet) Flush() error {
	return w.StreamWriter.Flush()
}

func (w *DynamicWriteSheet) FlushToFile() error {
	if len(w.SaveFilePath) == 0 {
		return errors.New("SaveFilePath is empty")
	}
	err := w.StreamWriter.Flush()
	if err != nil {
		return err
	}
	err = os.MkdirAll(filepath.Dir(w.SaveFilePath), os.ModePerm)
	if err != nil {
		return err
	}
	return w.File.SaveAs(w.SaveFilePath)
}

func (w *DynamicWriteSheet) Close() error {
	if w.File != nil {
		return w.File.Close()
	}
	return nil
}