158 lines
2.9 KiB
Go
158 lines
2.9 KiB
Go
package exporter
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/csv"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/xuri/excelize/v2"
|
|
)
|
|
|
|
type RowWriter interface {
|
|
WriteHeader(cols []string) error
|
|
WriteRow(vals []string) error
|
|
Close() (string, int64, error)
|
|
}
|
|
|
|
type CSVWriter struct {
|
|
f *os.File
|
|
buf *bufio.Writer
|
|
w *csv.Writer
|
|
count int64
|
|
}
|
|
|
|
func NewCSVWriter(dir, name string) (*CSVWriter, error) {
|
|
os.MkdirAll(dir, 0755)
|
|
p := filepath.Join(dir, name+"_"+time.Now().Format("20060102150405")+".csv")
|
|
f, err := os.Create(p)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
buf := bufio.NewWriterSize(f, 4<<20)
|
|
return &CSVWriter{f: f, buf: buf, w: csv.NewWriter(buf)}, nil
|
|
}
|
|
|
|
func (c *CSVWriter) WriteHeader(cols []string) error {
|
|
if err := c.w.Write(cols); err != nil {
|
|
return err
|
|
}
|
|
c.count++
|
|
return nil
|
|
}
|
|
|
|
func (c *CSVWriter) WriteRow(vals []string) error {
|
|
if err := c.w.Write(vals); err != nil {
|
|
return err
|
|
}
|
|
c.count++
|
|
return nil
|
|
}
|
|
|
|
func (c *CSVWriter) Close() (string, int64, error) {
|
|
c.w.Flush()
|
|
if c.buf != nil {
|
|
_ = c.buf.Flush()
|
|
}
|
|
p := c.f.Name()
|
|
info, _ := c.f.Stat()
|
|
c.f.Close()
|
|
return p, info.Size(), nil
|
|
}
|
|
|
|
type XLSXWriter struct {
|
|
f *excelize.File
|
|
sheet string
|
|
row int
|
|
path string
|
|
sw *excelize.StreamWriter
|
|
}
|
|
|
|
func NewXLSXWriter(dir, name, sheet string) (*XLSXWriter, error) {
|
|
os.MkdirAll(dir, 0755)
|
|
p := filepath.Join(dir, name+"_"+time.Now().Format("20060102150405")+".xlsx")
|
|
f := excelize.NewFile()
|
|
idx, err := f.GetSheetIndex(sheet)
|
|
if err != nil || idx < 0 {
|
|
idx, _ = f.NewSheet(sheet)
|
|
f.SetActiveSheet(idx)
|
|
if sheet != "Sheet1" {
|
|
_ = f.DeleteSheet("Sheet1")
|
|
}
|
|
} else {
|
|
f.SetActiveSheet(idx)
|
|
}
|
|
sw, e := f.NewStreamWriter(sheet)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
return &XLSXWriter{f: f, sheet: sheet, row: 1, path: p, sw: sw}, nil
|
|
}
|
|
|
|
func (x *XLSXWriter) WriteHeader(cols []string) error {
|
|
vals := make([]interface{}, len(cols))
|
|
for i := range cols {
|
|
vals[i] = cols[i]
|
|
}
|
|
axis := "A" + itoa(1)
|
|
if err := x.sw.SetRow(axis, vals); err != nil {
|
|
return err
|
|
}
|
|
x.row = 2
|
|
return nil
|
|
}
|
|
|
|
func (x *XLSXWriter) WriteRow(vals []string) error {
|
|
rowVals := make([]interface{}, len(vals))
|
|
for i := range vals {
|
|
rowVals[i] = vals[i]
|
|
}
|
|
axis := "A" + itoa(x.row)
|
|
if err := x.sw.SetRow(axis, rowVals); err != nil {
|
|
return err
|
|
}
|
|
x.row++
|
|
return nil
|
|
}
|
|
|
|
func (x *XLSXWriter) Close() (string, int64, error) {
|
|
if x.sw != nil {
|
|
_ = x.sw.Flush()
|
|
}
|
|
if err := x.f.SaveAs(x.path); err != nil {
|
|
return "", 0, err
|
|
}
|
|
info, err := os.Stat(x.path)
|
|
if err != nil {
|
|
return x.path, 0, nil
|
|
}
|
|
return x.path, info.Size(), nil
|
|
}
|
|
|
|
func col(n int) string {
|
|
s := ""
|
|
for n > 0 {
|
|
n--
|
|
s = string(rune('A'+(n%26))) + s
|
|
n /= 26
|
|
}
|
|
return s
|
|
}
|
|
|
|
func itoa(n int) string {
|
|
if n == 0 {
|
|
return "0"
|
|
}
|
|
b := make([]byte, 0, 10)
|
|
m := n
|
|
for m > 0 {
|
|
b = append(b, byte('0'+(m%10)))
|
|
m /= 10
|
|
}
|
|
for i, j := 0, len(b)-1; i < j; i, j = i+1, j-1 {
|
|
b[i], b[j] = b[j], b[i]
|
|
}
|
|
return string(b)
|
|
}
|