132 lines
2.7 KiB
Go
132 lines
2.7 KiB
Go
package exporter
|
|
|
|
import (
|
|
"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
|
|
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
|
|
}
|
|
return &CSVWriter{f: f, w: csv.NewWriter(f)}, 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()
|
|
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
|
|
}
|
|
|
|
func NewXLSXWriter(dir, name, sheet string) (*XLSXWriter, string, error) {
|
|
os.MkdirAll(dir, 0755)
|
|
p := filepath.Join(dir, name+"_"+time.Now().Format("20060102150405")+".xlsx")
|
|
f := excelize.NewFile()
|
|
f.NewSheet(sheet)
|
|
idx, _ := f.GetSheetIndex(sheet)
|
|
f.SetActiveSheet(idx)
|
|
return &XLSXWriter{f: f, sheet: sheet, row: 1}, p, nil
|
|
}
|
|
|
|
func (x *XLSXWriter) WriteHeader(cols []string) error {
|
|
for i, c := range cols {
|
|
cell := col(i+1) + "1"
|
|
if err := x.f.SetCellValue(x.sheet, cell, c); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
x.row = 2
|
|
return nil
|
|
}
|
|
|
|
func (x *XLSXWriter) WriteRow(vals []string) error {
|
|
r := x.row
|
|
for i, v := range vals {
|
|
cell := col(i+1) + itoa(r)
|
|
if err := x.f.SetCellValue(x.sheet, cell, v); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
x.row++
|
|
return nil
|
|
}
|
|
|
|
func (x *XLSXWriter) Close(path string) (string, int64, error) {
|
|
if err := x.f.SaveAs(path); err != nil {
|
|
return "", 0, err
|
|
}
|
|
info, err := os.Stat(path)
|
|
if err != nil {
|
|
return path, 0, nil
|
|
}
|
|
return path, info.Size(), nil
|
|
}
|
|
|
|
func col(n int) string {
|
|
s := ""
|
|
for n > 0 {
|
|
n--
|
|
s = string('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)
|
|
}
|