128 lines
1.9 KiB
Go
128 lines
1.9 KiB
Go
package export
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/csv"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"reflect"
|
|
)
|
|
|
|
type CSV struct {
|
|
fc *os.File
|
|
csv *csv.Writer
|
|
f *File
|
|
isNew bool
|
|
titles []string
|
|
}
|
|
|
|
func NewCSV(fileName string, limit int, param map[string]string) *CSV {
|
|
return &CSV{
|
|
f: NewFile(fileName, limit, param),
|
|
}
|
|
}
|
|
|
|
func (e *CSV) slice() {
|
|
if e.f.slice() {
|
|
e.reset()
|
|
}
|
|
}
|
|
|
|
func (e *CSV) SetParam(param map[string]string) {
|
|
e.f.param = param
|
|
}
|
|
|
|
func (e *CSV) reset() {
|
|
e.save()
|
|
e.f.NextFile()
|
|
e.Open()
|
|
e.WriteTitle(nil)
|
|
e.slice()
|
|
}
|
|
|
|
func (e *CSV) Open() error {
|
|
var err error
|
|
if e.f.IsFileExist() {
|
|
e.fc, err = os.OpenFile(e.f.FileName(), os.O_APPEND|os.O_WRONLY, 0644)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
e.f.SetRow(e.getLineCount(e.fc))
|
|
} else {
|
|
e.isNew = true
|
|
e.fc, err = os.Create(e.f.FileName())
|
|
}
|
|
|
|
if err == nil {
|
|
e.csv = csv.NewWriter(e.fc)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (e *CSV) getLineCount(file io.Reader) (line int) {
|
|
reader := bufio.NewReader(file)
|
|
line = 0
|
|
for {
|
|
_, isPrefix, err := reader.ReadLine()
|
|
if err != nil {
|
|
break
|
|
}
|
|
if !isPrefix {
|
|
line++
|
|
}
|
|
}
|
|
return line
|
|
}
|
|
|
|
func (e *CSV) save() error {
|
|
e.csv.Flush()
|
|
e.fc.Close()
|
|
return nil
|
|
}
|
|
|
|
func (e *CSV) WriteTitle(titles []string) error {
|
|
|
|
if titles != nil {
|
|
e.titles = titles
|
|
}
|
|
|
|
if e.titles != nil && e.isNew {
|
|
e.Write(e.titles)
|
|
e.isNew = false
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (e *CSV) Write(data interface{}) error {
|
|
if e.f.slice() {
|
|
e.reset()
|
|
}
|
|
|
|
v := reflect.ValueOf(data)
|
|
if v.Kind() == reflect.Ptr {
|
|
v = v.Elem()
|
|
}
|
|
if v.Kind() != reflect.Slice {
|
|
return errors.New("数据无效，不是切片类型")
|
|
}
|
|
|
|
switch val := data.(type) {
|
|
case []string:
|
|
return e.csv.Write(val)
|
|
case []interface{}:
|
|
strs := make([]string, len(val))
|
|
for i, v := range val {
|
|
strs[i] = fmt.Sprintf("%v", v)
|
|
}
|
|
return e.csv.Write(strs)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (e *CSV) Close() error {
|
|
return e.save()
|
|
}
|