bbgo_origin/pkg/bbgo/source.go

83 lines
2.2 KiB
Go
Raw Normal View History

package bbgo
import (
"encoding/json"
"strings"
"github.com/c9s/bbgo/pkg/fixedpoint"
"github.com/c9s/bbgo/pkg/types"
log "github.com/sirupsen/logrus"
)
type SourceFunc func(*types.KLine) fixedpoint.Value
type selectorInternal struct {
2022-09-06 10:15:13 +00:00
Source string
sourceGetter SourceFunc
}
func (s *selectorInternal) UnmarshalJSON(d []byte) error {
if err := json.Unmarshal(d, &s.Source); err != nil {
return err
}
s.init()
return nil
}
func (s selectorInternal) MarshalJSON() ([]byte, error) {
if s.Source == "" {
s.Source = "close"
s.init()
}
return []byte("\"" + s.Source + "\""), nil
}
type SourceSelector struct {
Source selectorInternal `json:"source,omitempty"`
}
func (s *selectorInternal) init() {
switch strings.ToLower(s.Source) {
case "close":
2022-09-06 10:15:13 +00:00
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value { return kline.Close }
case "high":
2022-09-06 10:15:13 +00:00
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value { return kline.High }
case "low":
2022-09-06 10:15:13 +00:00
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value { return kline.Low }
case "hl2":
2022-09-06 10:15:13 +00:00
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value { return kline.High.Add(kline.Low).Div(fixedpoint.Two) }
case "hlc3":
2022-09-06 10:15:13 +00:00
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value {
return kline.High.Add(kline.Low).Add(kline.Close).Div(fixedpoint.Three)
}
case "ohlc4":
2022-09-06 10:15:13 +00:00
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value {
return kline.High.Add(kline.Low).Add(kline.Close).Add(kline.Open).Div(fixedpoint.Four)
}
case "open":
2022-09-06 10:15:13 +00:00
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value { return kline.Open }
case "oc2":
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value { return kline.Open.Add(kline.Close).Div(fixedpoint.Two) }
default:
log.Infof("source not set: %s, use hl2 by default", s.Source)
2022-09-06 10:15:13 +00:00
s.sourceGetter = func(kline *types.KLine) fixedpoint.Value { return kline.High.Add(kline.Low).Div(fixedpoint.Two) }
}
}
func (s *selectorInternal) String() string {
if s.Source == "" {
s.Source = "close"
s.init()
}
return s.Source
}
// lazy init if empty struct is passed in
func (s *SourceSelector) GetSource(kline *types.KLine) fixedpoint.Value {
if s.Source.Source == "" {
s.Source.Source = "close"
s.Source.init()
}
2022-09-06 10:15:13 +00:00
return s.Source.sourceGetter(kline)
}