move account type into types package

pkg/bbgo/account.go

@@ -1,47 +1,2 @@
 package bbgo
 
-import (
-	"sync"
-
-	"github.com/c9s/bbgo/pkg/types"
-	"github.com/c9s/bbgo/pkg/util"
-
-	log "github.com/sirupsen/logrus"
-)
-
-type Account struct {
-	sync.Mutex
-	balances map[string]types.Balance
-}
-
-func (a *Account) Balance(currency string) (balance types.Balance, ok bool) {
-	a.Lock()
-	balance, ok = a.balances[currency]
-	a.Unlock()
-	return balance, ok
-}
-
-func (a *Account) handleBalanceUpdates(balances map[string]types.Balance) {
-	a.Lock()
-	defer a.Unlock()
-
-	for _, balance := range balances {
-		a.balances[balance.Currency] = balance
-	}
-}
-
-func (a *Account) BindStream(stream types.Stream) {
-	stream.OnBalanceUpdate(a.handleBalanceUpdates)
-	stream.OnBalanceSnapshot(a.handleBalanceUpdates)
-}
-
-func (a *Account) Print() {
-	a.Lock()
-	defer a.Unlock()
-
-	for _, balance := range a.balances {
-		if util.NotZero(balance.Available) {
-			log.Infof("account balance %s %f", balance.Currency, balance.Available)
-		}
-	}
-}

pkg/bbgo/accounting.go

@@ -1 +0,0 @@
-package bbgo

pkg/bbgo/environment.go

@@ -112,7 +112,8 @@ func (environ *Environment) Init(ctx context.Context) (err error) {
 
 		stream := session.Exchange.NewStream()
 
-		account := &Account{balances: balances}
+		account := &types.Account{}
+		account.UpdateBalances(balances)
 		account.BindStream(stream)
 		session.Account = account
 

pkg/bbgo/session.go

@@ -9,7 +9,7 @@ type ExchangeSession struct {
 	Name string
 
 	// The exchange account states
-	Account *Account
+	Account *types.Account
 
 	// Stream is the connection stream of the exchange
 	Stream types.Stream

pkg/exchange/binance/exchange.go

@@ -237,7 +237,7 @@ func (e *Exchange) QueryAccountBalances(ctx context.Context) (types.BalanceMap,
 		return nil, err
 	}
 
-	return account.Balances, nil
+	return account.Balances(), nil
 }
 
 // PlatformFeeCurrency
@@ -260,11 +260,12 @@ func (e *Exchange) QueryAccount(ctx context.Context) (*types.Account, error) {
 		}
 	}
 
-	return &types.Account{
+	a := &types.Account{
 		MakerCommission: account.MakerCommission,
 		TakerCommission: account.TakerCommission,
-		Balances:        balances,
-	}, nil
+	}
+	a.UpdateBalances(balances)
+	return a, nil
 }
 
 func (e *Exchange) SubmitOrder(ctx context.Context, order types.SubmitOrder) error {

pkg/exchange/max/exchange.go

@@ -114,11 +114,13 @@ func (e *Exchange) QueryAccount(ctx context.Context) (*types.Account, error) {
 		}
 	}
 
-	return &types.Account{
+	a := &types.Account{
 		MakerCommission: 15, // 0.15%
 		TakerCommission: 15, // 0.15%
-		Balances:        balances,
-	}, nil
+	}
+
+	a.UpdateBalances(balances)
+	return a, nil
 }
 
 func (e *Exchange) QueryWithdrawHistory(ctx context.Context, asset string, since, until time.Time) (allWithdraws []types.Withdraw, err error) {

pkg/types/account.go

@@ -1,5 +1,13 @@
 package types
 
+import (
+	"sync"
+
+	"github.com/sirupsen/logrus"
+
+	"github.com/c9s/bbgo/pkg/util"
+)
+
 type Balance struct {
 	Currency  string  `json:"currency"`
 	Available float64 `json:"available"`
@@ -9,18 +17,59 @@ type Balance struct {
 type BalanceMap map[string]Balance
 
 type Account struct {
+	sync.Mutex
+
 	MakerCommission int64
 	TakerCommission int64
 	AccountType     string
-	Balances        map[string]Balance
+
+	balances BalanceMap
 }
 
-func (a *Account) UpdateBalance(b Balance) {
-	a.Balances[b.Currency] = b
+func (a *Account) Balances() BalanceMap {
+	d := make(BalanceMap)
+
+	a.Lock()
+	for c, b := range a.balances {
+		d[c] = b
+	}
+	a.Unlock()
+
+	return d
 }
 
-func NewAccount() *Account {
-	return &Account{
-		Balances: make(map[string]Balance),
+func (a *Account) Balance(currency string) (balance Balance, ok bool) {
+	a.Lock()
+	balance, ok = a.balances[currency]
+	a.Unlock()
+	return balance, ok
+}
+
+func (a *Account) UpdateBalances(balances map[string]Balance) {
+	a.Lock()
+	defer a.Unlock()
+
+	if a.balances == nil {
+		a.balances = make(BalanceMap)
+	}
+
+	for _, balance := range balances {
+		a.balances[balance.Currency] = balance
+	}
+}
+
+func (a *Account) BindStream(stream Stream) {
+	stream.OnBalanceUpdate(a.UpdateBalances)
+	stream.OnBalanceSnapshot(a.UpdateBalances)
+}
+
+func (a *Account) Print() {
+	a.Lock()
+	defer a.Unlock()
+
+	for _, balance := range a.balances {
+		if util.NotZero(balance.Available) {
+			logrus.Infof("account balance %s %f", balance.Currency, balance.Available)
+		}
 	}
 }