diff --git a/pkg/core/tradecollector.go b/pkg/core/tradecollector.go index e6e5d515f..da3e19445 100644 --- a/pkg/core/tradecollector.go +++ b/pkg/core/tradecollector.go @@ -12,6 +12,8 @@ import ( "github.com/c9s/bbgo/pkg/types" ) +type TradeConverter func(trade types.Trade) (types.Trade, error) + //go:generate callbackgen -type TradeCollector type TradeCollector struct { Symbol string @@ -25,6 +27,8 @@ type TradeCollector struct { mu sync.Mutex + tradeConverters []TradeConverter + recoverCallbacks []func(trade types.Trade) tradeCallbacks []func(trade types.Trade, profit, netProfit fixedpoint.Value) @@ -49,6 +53,24 @@ func NewTradeCollector(symbol string, position *types.Position, orderStore *Orde } } +func (c *TradeCollector) AddTradeConverter(converter TradeConverter) { + c.tradeConverters = append(c.tradeConverters, converter) +} + +func (c *TradeCollector) convertTrade(trade types.Trade) types.Trade { + for _, converter := range c.tradeConverters { + convTrade, err := converter(trade) + if err != nil { + logrus.WithError(err).Errorf("trade converter error, trade: %s", trade.String()) + continue + } + + trade = convTrade + } + + return trade +} + // OrderStore returns the order store used by the trade collector func (c *TradeCollector) OrderStore() *OrderStore { return c.orderStore