gocommon/pkg/common/excel.go

package common

import (
	"errors"
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/360EntSecGroup-Skylar/excelize/v2"
)

// #################
//   excel 对外接口 , demo 可以见test
// #################
type Excel interface {
	WriteSheet(sheetName string, sheetHeader []string, writer RowWriter) error // 写入sheet
	FlushWrite(writer io.Writer) error                                         // 写出去
	GetSheetList() []string                                                    // 获取sheet-list
	ReadSheet(sheetName string) ([][]string, error)                            // 读取sheet
}

type Row interface {
	Write(v interface{}) // 每写入一个数据，ColumnIndex()就会+1, 推荐v为字符串类型
	ColumnIndex() int    // 获取当前列所在的index，从0开始 (索引从0开始，主要是为了和数组对应起来)
	RowIndex() int       // 获取当前行所在的index，从0开始 (索引从0开始，主要是为了和数组对应起来)
}

type RowWriter func(r Row) (hasNext bool) // 是否还要写入下一行

// #################
//   360 excel 实现
// #################
type _360Excel struct {
	*excelize.File
}

func NewWriter360Excel() Excel {
	file := excelize.NewFile()
	return &_360Excel{
		File: file,
	}
}

func NewReader360Excel(reader io.Reader) (Excel, error) {
	excel := new(_360Excel)
	r, err := excelize.OpenReader(reader)
	if err != nil {
		return nil, err
	}
	excel.File = r
	return excel, nil
}

type row struct {
	*excelize.File        // out
	columnNum      int    // 第几列,0开始
	rowNum         int    // 第几行,excel的行数,从1开始，如果写了header从2开始
	offset         int    // 行的偏移量
	sheetName      string // sheet_name
	err            error  // 运行中间的异常
	builder        strings.Builder
}

func newRow(f *excelize.File, sheetName string, startRowNum int) *row {
	return &row{
		File:      f,
		sheetName: sheetName,
		rowNum:    startRowNum,
		offset:    startRowNum, // row偏差
		columnNum: 1,           // 初始化为1
	}
}

func (r *row) Write(v interface{}) {
	if r.err != nil {
		return
	}
	columnName, err := excelize.ColumnNumberToName(r.columnNum)
	if err != nil {
		r.err = err
		return
	}

	r.builder.WriteString(columnName)
	r.builder.WriteString(rowNumberName(r.rowNum))
	if err := r.SetCellValue(r.sheetName, r.builder.String(), v); err != nil {
		r.err = err
		return
	}

	// end handler
	r.builder.Reset()
	r.columnNum++
}

func (r *row) ColumnIndex() int {
	return r.columnNum - 1 // 偏差
}

func (r *row) RowIndex() int {
	return r.rowNum - r.offset // 偏差
}

func (r *row) increaseRow() {
	r.rowNum++
	r.columnNum = 1
}

func (e *_360Excel) FlushWrite(writer io.Writer) error {
	e.DeleteSheet("Sheet1")
	e.SetActiveSheet(0)
	return e.Write(writer)
}

func rowNumberName(rowNum int) string {
	return strconv.Itoa(rowNum)
}

// excel的index都是从1开始,和传统计数都有偏差！
func (e *_360Excel) WriteSheet(sheetName string, sheetHeader []string, writer RowWriter) (err error) {
	// 抓取recover
	defer func() {
		if rerr := recover(); rerr != nil {
			err = errors.New(fmt.Sprintf("%+v", rerr))
		}
	}()

	// 校验，默认会有一个sheet1，所以不能和他使用的一样，因此需要最后删除
	if sheetName == "" || sheetName == "Sheet1" {
		return errors.New("sheetName不能为空或为Sheet1")
	}
	// 初始化 sheet
	e.NewSheet(sheetName)

	var (
		rowNum    = 1                                          // excel的行是从1开始计数
		hasHeader = sheetHeader != nil && len(sheetHeader) > 0 // 是否有header
	)
	// header A1 B1 C1 D1 ... 写header
	for index, title := range sheetHeader {
		columnNumberName, err := excelize.ColumnNumberToName(index + 1) // excel列从A开始
		if err != nil {
			return err
		}
		if err := e.SetCellStr(sheetName, columnNumberName+rowNumberName(rowNum), title); err != nil {
			return err
		}
	}
	// row ++
	if hasHeader {
		rowNum++
	}

	row := newRow(e.File, sheetName, rowNum)
	// 循环写
	for writer(row) {
		if err := row.err; err != nil {
			return err
		}
		row.increaseRow()
	}
	return nil
}

func (e *_360Excel) GetSheetList() []string {
	return e.File.GetSheetList()
}

func (e *_360Excel) ReadSheet(sheetName string) ([][]string, error) {
	result, err := e.GetRows(sheetName)
	if err != nil {
		return nil, err
	}
	return result, nil
}