|
|
package csv
|
|
|
|
|
|
import (
|
|
|
"context"
|
|
|
"joybuywork/internal/plugins/export"
|
|
|
"encoding/csv"
|
|
|
"fmt"
|
|
|
"os"
|
|
|
)
|
|
|
|
|
|
type Export struct {
|
|
|
export.BaseExport
|
|
|
fields []string //表头对应的字段key []string{"key1", "key2"}
|
|
|
header []string //表头数据 []string{"表头1", “表头2”}
|
|
|
file *os.File //文件资源
|
|
|
w *csv.Writer //csv文件写入资源
|
|
|
}
|
|
|
|
|
|
// NewExport 实例化csv导出对象
|
|
|
func NewExport(ctx context.Context, header export.HeaderField, total int64, opts ...export.Option) *Export {
|
|
|
ept := new(Export)
|
|
|
{
|
|
|
ept.Init(ctx, header, total, opts...)
|
|
|
}
|
|
|
|
|
|
//初始化头部信息和字段key信息
|
|
|
for _, v := range ept.Header {
|
|
|
ept.header = append(ept.header, v.Title)
|
|
|
ept.fields = append(ept.fields, v.Field)
|
|
|
}
|
|
|
|
|
|
//设置文件名称
|
|
|
ept.SetFilename(ept)
|
|
|
|
|
|
return ept
|
|
|
}
|
|
|
|
|
|
// ExportFile 直接导出数据
|
|
|
func (e *Export) ExportFile(data []export.RowData) (err error) {
|
|
|
defer e.file.Close()
|
|
|
|
|
|
err = e.initFile()
|
|
|
if err != nil {
|
|
|
return
|
|
|
}
|
|
|
|
|
|
list := e.getList(data)
|
|
|
|
|
|
//写入数据
|
|
|
err = e.w.WriteAll(list)
|
|
|
if err != nil {
|
|
|
return err
|
|
|
}
|
|
|
|
|
|
return
|
|
|
}
|
|
|
|
|
|
// 初始化csv文件对象
|
|
|
func (e *Export) initFile() (err error) {
|
|
|
// 不存在则创建;存在则追加;读写模式;
|
|
|
file, err := os.OpenFile(e.GetFullPath(), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
|
|
|
if err != nil {
|
|
|
return
|
|
|
}
|
|
|
|
|
|
//自定义文件
|
|
|
if e.IsCustomerFile() {
|
|
|
e.w = csv.NewWriter(file)
|
|
|
return
|
|
|
}
|
|
|
|
|
|
// 写入UTF-8 BOM，防止中文乱码
|
|
|
_, err = file.WriteString("\xEF\xBB\xBF")
|
|
|
if err != nil {
|
|
|
return
|
|
|
}
|
|
|
|
|
|
//实例csv写入数据
|
|
|
w := csv.NewWriter(file)
|
|
|
|
|
|
//写把头写到数据里面
|
|
|
err = w.Write(e.header)
|
|
|
if err != nil {
|
|
|
return
|
|
|
}
|
|
|
w.Flush()
|
|
|
|
|
|
e.w = w
|
|
|
e.file = file
|
|
|
|
|
|
return
|
|
|
}
|
|
|
|
|
|
// 获取每行输出的数据
|
|
|
func (e *Export) getList(data []export.RowData) [][]string {
|
|
|
var list [][]string
|
|
|
for _, row := range data {
|
|
|
var one []string
|
|
|
for _, field := range e.fields {
|
|
|
v, ok := row[field]
|
|
|
if ok {
|
|
|
one = append(one, fmt.Sprintf("%v", v))
|
|
|
} else {
|
|
|
one = append(one, "")
|
|
|
}
|
|
|
}
|
|
|
list = append(list, one)
|
|
|
}
|
|
|
|
|
|
return list
|
|
|
}
|
|
|
|
|
|
// FileSuffix 文件后缀
|
|
|
func (e *Export) FileSuffix() string {
|
|
|
return "csv"
|
|
|
}
|