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) }