fix kline dumper

This commit is contained in:
c9s 2022-05-16 22:24:15 +08:00
parent 18b2ea077c
commit e651b9d36f
No known key found for this signature in database
GPG Key ID: 7385E7E464CB0A54

View File

@ -1,15 +1,14 @@
package backtest package backtest
import ( import (
"encoding/csv"
"fmt" "fmt"
"os"
"path/filepath" "path/filepath"
"strconv" "strconv"
"time" "time"
"go.uber.org/multierr" "go.uber.org/multierr"
"github.com/c9s/bbgo/pkg/data/tsv"
"github.com/c9s/bbgo/pkg/types" "github.com/c9s/bbgo/pkg/types"
) )
@ -23,16 +22,14 @@ type symbolInterval struct {
// KLineDumper dumps the received kline data into a folder for the backtest report to load the charts. // KLineDumper dumps the received kline data into a folder for the backtest report to load the charts.
type KLineDumper struct { type KLineDumper struct {
OutputDirectory string OutputDirectory string
files map[symbolInterval]*os.File writers map[symbolInterval]*tsv.Writer
writers map[symbolInterval]*csv.Writer
filenames map[symbolInterval]string filenames map[symbolInterval]string
} }
func NewKLineDumper(outputDirectory string) *KLineDumper { func NewKLineDumper(outputDirectory string) *KLineDumper {
return &KLineDumper{ return &KLineDumper{
OutputDirectory: outputDirectory, OutputDirectory: outputDirectory,
files: make(map[symbolInterval]*os.File), writers: make(map[symbolInterval]*tsv.Writer),
writers: make(map[symbolInterval]*csv.Writer),
filenames: make(map[symbolInterval]string), filenames: make(map[symbolInterval]string),
} }
} }
@ -42,7 +39,7 @@ func (d *KLineDumper) Filenames() map[symbolInterval]string {
} }
func (d *KLineDumper) formatFileName(symbol string, interval types.Interval) string { func (d *KLineDumper) formatFileName(symbol string, interval types.Interval) string {
return filepath.Join(d.OutputDirectory, fmt.Sprintf("%s-%s.csv", return filepath.Join(d.OutputDirectory, fmt.Sprintf("%s-%s.tsv",
symbol, symbol,
interval)) interval))
} }
@ -69,36 +66,28 @@ func (d *KLineDumper) Record(k types.KLine) error {
w, ok := d.writers[si] w, ok := d.writers[si]
if !ok { if !ok {
filename := d.formatFileName(k.Symbol, k.Interval) filename := d.formatFileName(k.Symbol, k.Interval)
f2, err := os.Create(filename) w2, err := tsv.NewWriterFile(filename)
if err != nil { if err != nil {
return err return err
} }
w = w2
w = csv.NewWriter(f2) d.writers[si] = w2
d.files[si] = f2
d.writers[si] = w
d.filenames[si] = filename d.filenames[si] = filename
if err := w.Write(csvHeader); err != nil { if err2 := w2.Write(csvHeader); err2 != nil {
return err return err2
} }
} }
if err := w.Write(d.encode(k)); err != nil { return w.Write(d.encode(k))
return err
}
return nil
} }
func (d *KLineDumper) Close() error { func (d *KLineDumper) Close() error {
var err error = nil
for _, w := range d.writers { for _, w := range d.writers {
w.Flush() w.Flush()
} err2 := w.Close()
var err error = nil
for _, f := range d.files {
err2 := f.Close()
if err2 != nil { if err2 != nil {
err = multierr.Append(err, err2) err = multierr.Append(err, err2)
} }