joybuywork/internal/plugins/export/xlsx/xlsx.go

package xlsx

import (
	"context"
	"joybuywork/internal/plugins/export"
	"fmt"

	"github.com/xuri/excelize/v2"
)

const (
	DefaultSheetName = "Sheet1" //默认拥有的工作簿名称
)

const (
	defaultSfIdx = 2 //每个sheet写入行从第二行开始
)

var columnField = []string{"A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"}

type Export struct {
	export.BaseExport
	file *excelize.File //xlsx的文件对象

	//多sheet模式需要的参数
	ss    []Sheet        //要导出的sheet
	ssIdx int            //当前导出的sheet索引
	swIdx int            //当前切换的idx
	sfv   map[string]int //切换字段值对应的sheet的索引（PS：当设定了sheet对应字段是对应的值得时候，就会写到那个sheet上）

	//下面都是不用多sheet模式需要的参数
	sfIdx  int64             //当前sheet写到第几行
	sheet  string            //当前写入的sheet
	fields map[string]string //头部对应的表格列和字段名 map[string]string{"A":"字段key1", "B": "字段key2"}
	header map[string]string //表头数据map[string]string{"A1":"表头1", "B1": "表头2"}
}

// NewExport 实例化一个xlsx导出对象
func NewExport(ctx context.Context, header export.HeaderField, total int64, opts ...export.Option) *Export {
	ept := new(Export)
	{
		ept.Init(ctx, header, total, opts...)
	}

	//初始化表头信息
	ept.header, ept.fields = ept.getHeaderAndFields()

	//设置文件名称
	ept.SetFilename(ept)

	//初始化文件
	ept.initFile()

	//默认的工作簿名称
	ept.sheet, ept.sfIdx = DefaultSheetName, defaultSfIdx

	return ept
}

// Use 使用一些参数
func (e *Export) Use(opts ...Option) *Export {
	for _, opt := range opts {
		opt(e)
	}
	return e
}

// 初始化文件
func (e *Export) initFile() {
	if e.IsCustomerFile() {
		e.file, _ = excelize.OpenFile(e.GetFullPath())
	} else {
		e.file = excelize.NewFile()
	}
}

// 自动切换sheet
func (e *Export) checkSwitchSheet() (ok bool) {
	if !e.HasManySheet() || e.sfv != nil {
		return
	}

	//当前sheet是不是写满了，写满了就切
	if e.ss[e.ssIdx].writeNum == e.ss[e.ssIdx].Total {
		ok = e.switchSheet()
	} else {
		//当前sheet写入多一条数据
		e.ss[e.ssIdx].writeNum++
	}

	return
}

// 切换sheet
func (e *Export) switchSheet() (sws bool) {
	l := len(e.ss)

	//已经是最后一页
	if l-1 == e.swIdx {
		return
	}

	//剩余的sheet
	remain := e.ss[e.swIdx+1:]
	for _, v := range remain {
		//下一个sheet
		e.swIdx++
		if v.Total > 0 {
			sws = true
			break
		}
	}

	//有下一个sheet
	if sws {
		//把切换成功的idx换成导出的idx
		e.ssIdx, e.sfIdx = e.swIdx, defaultSfIdx

		//初始化sheet
		e.initSheet(&e.ss[e.ssIdx], false)
	}

	return
}

// 初始化sheet
func (e *Export) initSheet(sheet *Sheet, setHead bool) {
	e.sheet, e.Header = sheet.Name, sheet.Header

	if sheet.header == nil || sheet.fields == nil {
		//初始化字段数据
		sheet.header, sheet.fields = e.getHeaderAndFields()
		sheet.sfIdx = defaultSfIdx
	}

	//设置当前导出sheet
	e.header, e.fields = sheet.header, sheet.fields

	//设置sheet的表头
	if setHead {
		_ = e.setHeader()
	}
}

// 获取字段信息字段信息
func (e *Export) getHeaderAndFields() (map[string]string, map[string]string) {
	//初始化导出表头和字段数据
	header, fields := make(map[string]string), make(map[string]string)
	for k, v := range e.Header {
		cf := e.getColumnField(k)
		fields[cf] = v.Field
		header[fmt.Sprintf("%s1", cf)] = v.Title //第一行写表头
	}
	return header, fields
}

// 获取column
func (e *Export) getColumnField(idx int) string {
	if idx < len(columnField) {
		return columnField[idx]
	}

	f := (idx / len(columnField)) - 1
	i := idx % len(columnField)

	return fmt.Sprintf("%s%s", columnField[f], columnField[i])

}

type SheetRowData struct {
	sheet string         //当前行写入的sheet
	row   map[string]any //当前行的数据
}

// 获取写入xlsx的数据行
func (e *Export) getList(data []export.RowData) (list []SheetRowData) {
	list = make([]SheetRowData, 0, len(data))
	for _, row := range data {
		list = append(list, e.getSheetData(row))
	}
	return
}

// 获取sheet的数据
func (e *Export) getSheetData(row export.RowData) SheetRowData {
	fields, idx, sfIdx := e.fields, e.ssIdx, e.sfIdx

	var one = SheetRowData{
		sheet: e.sheet,
		row:   map[string]any{},
	}

	//如果事标识的分sheet
	if e.sfv != nil {
		for field, value := range row {
			if i, ok := e.sfv[e.getSfvKey(field, value)]; ok {
				idx = i
				fields, one.sheet, sfIdx = e.ss[idx].fields, e.ss[idx].Name, e.ss[idx].sfIdx
			}
		}
		e.ss[idx].sfIdx++
	}

	//当前行的数据
	for column, field := range fields {
		columnKey := fmt.Sprintf("%s%d", column, sfIdx)
		one.row[columnKey] = ""
		if v, ok := row[field]; ok {
			one.row[columnKey] = v
		}
	}
	e.sfIdx++

	return one
}

// 初始化xlsx文件对象
func (e *Export) setHeader() (err error) {
	for axis, value := range e.header {
		err = e.file.SetCellValue(e.sheet, axis, value)
		if err != nil {
			return
		}
	}
	return
}

// DeleteSheet 删除工作表
func (e *Export) DeleteSheet(name string) {
	e.file.DeleteSheet(name)
}

// CreateSheet 创建sheet
func (e *Export) CreateSheet(name string) (idx int) {
	idx, _ = e.file.GetSheetIndex(name)
	if idx == -1 {
		idx, _ = e.file.NewSheet(name)
	}
	return
}

// ExportFile 直接导出数据
func (e *Export) ExportFile(data []export.RowData) (err error) {
	err = e.setHeader()
	if err != nil {
		return
	}

	//工作簿写入
	e.sheetWrite(data)

	return e.file.SaveAs(e.GetFullPath())
}

// 工作簿写入
func (e *Export) sheetWrite(data []export.RowData) {
	list := e.getList(data)
	for key, row := range list {
		//有标识切sheet的时候
		ok := e.checkSwitchSheet()
		if ok { //切换成功，要重新计算行数
			e.sheetWrite(data[key:])
			return
		}

		//写入文件
		for axis, value := range row.row {
			switch value.(type) {
			case int:
				_ = e.file.SetCellInt(row.sheet, axis, value.(int))
			case int32:
				_ = e.file.SetCellInt(row.sheet, axis, int(value.(int32)))
			case uint32:
				_ = e.file.SetCellInt(row.sheet, axis, int(value.(uint32)))
			case int64:
				_ = e.file.SetCellInt(row.sheet, axis, int(value.(int64)))
			case uint64:
				_ = e.file.SetCellInt(row.sheet, axis, int(value.(uint64)))
			default:
				_ = e.file.SetCellValue(row.sheet, axis, fmt.Sprintf("%v\t", value))
			}
		}
	}
}

// FileSuffix 文件后缀
func (e *Export) FileSuffix() string {
	return "xlsx"
}

// HasManySheet 是否有多sheet导出
func (e *Export) HasManySheet() bool {
	return len(e.ss) > 0
}

// 获取sfv的key
func (e *Export) getSfvKey(key string, val interface{}) string {
	return fmt.Sprintf("%s_%v", key, val)
}