add balance type

config/rebalance.yaml

@@ -12,5 +12,6 @@ exchangeStrategies:
       maxAmount: 1_000 # max amount to buy or sell per order
       orderType: LIMIT_MAKER # LIMIT, LIMIT_MAKER or MARKET
       priceType: MAKER # LAST, MID, TAKER or MAKER
+      balanceType: TOTAL
       dryRun: true
       onStart: true

pkg/strategy/rebalance/strategy.go

@@ -34,6 +34,7 @@ type Strategy struct {
 	MaxAmount     fixedpoint.Value  `json:"maxAmount"` // max amount to buy or sell per order
 	OrderType     types.OrderType   `json:"orderType"`
 	PriceType     types.PriceType   `json:"priceType"`
+	BalanceType   types.BalanceType `json:"balanceType"`
 	DryRun        bool              `json:"dryRun"`
 	OnStart       bool              `json:"onStart"` // rebalance on start
 
@@ -51,6 +52,10 @@ func (s *Strategy) Defaults() error {
 	if s.PriceType == "" {
 		s.PriceType = types.PriceTypeMaker
 	}
+
+	if s.BalanceType == "" {
+		s.BalanceType = types.BalanceTypeAvailable
+	}
 	return nil
 }
 
@@ -216,7 +221,7 @@ func (s *Strategy) generateOrder(ctx context.Context) (*types.SubmitOrder, error
 		return nil, err
 	}
 
-	values := prices.Mul(toValueMap(balances))
+	values := prices.Mul(s.toValueMap(balances))
 	weights := values.Normalize()
 
 	for symbol, market := range s.markets {
@@ -286,11 +291,10 @@ func (s *Strategy) generateOrder(ctx context.Context) (*types.SubmitOrder, error
 	return nil, nil
 }
 
-func toValueMap(balances types.BalanceMap) types.ValueMap {
+func (s *Strategy) toValueMap(balances types.BalanceMap) types.ValueMap {
 	m := make(types.ValueMap)
 	for _, b := range balances {
-		// m[b.Currency] = b.Net()
-		m[b.Currency] = b.Available
+		m[b.Currency] = s.BalanceType.Map(b)
 	}
 	return m
 }

pkg/types/balance_type.go

@@ -0,0 +1,69 @@
+package types
+
+import (
+	"encoding/json"
+	"strings"
+
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+	"github.com/pkg/errors"
+)
+
+type BalanceType string
+
+const (
+	BalanceTypeAvailable BalanceType = "AVAILABLE"
+	BalanceTypeLocked    BalanceType = "LOCKED"
+	BalanceTypeBorrowed  BalanceType = "BORROWED"
+	BalanceTypeInterest  BalanceType = "INTEREST"
+	BalanceTypeNet       BalanceType = "NET"
+	BalanceTypeTotal     BalanceType = "TOTAL"
+	BalanceTypeDebt      BalanceType = "DEBT"
+)
+
+var ErrInvalidBalanceType = errors.New("invalid balance type")
+
+func ParseBalanceType(s string) (b BalanceType, err error) {
+	b = BalanceType(strings.ToUpper(s))
+	switch b {
+	case BalanceTypeAvailable, BalanceTypeLocked, BalanceTypeBorrowed, BalanceTypeInterest, BalanceTypeNet, BalanceTypeTotal, BalanceTypeDebt:
+		return b, nil
+	}
+	return b, ErrInvalidBalanceType
+}
+
+func (b *BalanceType) UnmarshalJSON(data []byte) error {
+	var s string
+
+	if err := json.Unmarshal(data, &s); err != nil {
+		return err
+	}
+
+	t, err := ParseBalanceType(s)
+	if err != nil {
+		return err
+	}
+
+	*b = t
+
+	return nil
+}
+
+func (b BalanceType) Map(balance Balance) fixedpoint.Value {
+	switch b {
+	case BalanceTypeAvailable:
+		return balance.Available
+	case BalanceTypeLocked:
+		return balance.Locked
+	case BalanceTypeBorrowed:
+		return balance.Borrowed
+	case BalanceTypeInterest:
+		return balance.Interest
+	case BalanceTypeNet:
+		return balance.Net()
+	case BalanceTypeTotal:
+		return balance.Total()
+	case BalanceTypeDebt:
+		return balance.Debt()
+	}
+	return fixedpoint.Zero
+}

pkg/types/price_type.go

@@ -21,24 +21,13 @@ const (
 
 var ErrInvalidPriceType = errors.New("invalid price type")
 
-func StrToPriceType(s string) (price PriceType, err error) {
-	switch strings.ToLower(s) {
-	case "last":
-		price = PriceTypeLast
-	case "buy":
-		price = PriceTypeBuy
-	case "sell":
-		price = PriceTypeSell
-	case "mid":
-		price = PriceTypeMid
-	case "maker":
-		price = PriceTypeMaker
-	case "taker":
-		price = PriceTypeTaker
-	default:
-		err = ErrInvalidPriceType
+func ParsePriceType(s string) (p PriceType, err error) {
+	p = PriceType(strings.ToUpper(s))
+	switch p {
+	case PriceTypeLast, PriceTypeBuy, PriceTypeSell, PriceTypeMid, PriceTypeMaker, PriceTypeTaker:
+		return p, err
 	}
-	return price, err
+	return p, ErrInvalidPriceType
 }
 
 func (p *PriceType) UnmarshalJSON(data []byte) error {
@@ -48,12 +37,13 @@ func (p *PriceType) UnmarshalJSON(data []byte) error {
 		return err
 	}
 
-	t, err := StrToPriceType(s)
+	t, err := ParsePriceType(s)
 	if err != nil {
 		return err
 	}
 
 	*p = t
+
 	return nil
 }