python: fix market service

python/bbgo/data/__init__.py

@@ -1,4 +1,8 @@
 from .balance import Balance
+from .depth import Depth
+from .depth import PriceVolume
 from .error import ErrorMessage
+from .event import Event
 from .kline import KLine
 from .order import Order
+from .subscription import Subscription

python/bbgo/data/depth.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import bbgo_pb2
+
+from typing import List
+
+
+# message Depth {
+#   string exchange = 1;
+#   string symbol = 2;
+#   repeated PriceVolume asks = 3;
+#   repeated PriceVolume bids = 4;
+# }
+@dataclass
+class Depth:
+    exchange: str
+    symbol: str
+    asks: List[PriceVolume]
+    bids: List[PriceVolume]
+
+    @classmethod
+    def from_pb(cls, obj: bbgo_pb2.Depth):
+        return cls(
+            exchange=obj.exchange,
+            symbol=obj.symbol,
+            asks=[PriceVolume.from_pb(ask) for ask in obj.asks],
+            bids=[PriceVolume.from_pb(bid) for bid in obj.bids],
+        )
+
+
+@dataclass
+class PriceVolume:
+    price: float
+    volume: float
+
+    @classmethod
+    def from_pb(cls, obj: bbgo_pb2.PriceVolume):
+        return cls(
+            price=float(obj.price),
+            volume=float(obj.volume),
+        )

python/bbgo/data/event.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import bbgo_pb2
+
+from ..enums import EventType, ChannelType
+from . import Depth
+
+
+@dataclass
+class Event:
+    exchange: str
+    symbol: str
+    channel_type: ChannelType
+    event_type: EventType
+    depth: Depth = None
+
+    @classmethod
+    def from_pb(cls, obj: bbgo_pb2.SubscribeResponse) -> Event:
+        channel_type = ChannelType.from_pb(obj.channel)
+
+        event = cls(
+            exchange=obj.exchange,
+            symbol=obj.symbol,
+            channel_type=channel_type,
+            event_type=EventType.from_pb(obj.event),
+        )
+
+        if channel_type == ChannelType.BOOK:
+            event.depth = Depth.from_pb(obj.depth)
+
+        return event

python/bbgo/data/subscription.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import bbgo_pb2
+
+from ..enums import ChannelType
+from ..enums import DepthType
+
+
+@dataclass
+class Subscription:
+    exchange: str
+    channel: ChannelType
+    symbol: str
+    depth: DepthType = None
+    interval: str = None
+
+    def to_pb(self) -> bbgo_pb2.Subscription:
+        subscription_pb = bbgo_pb2.Subscription(
+            exchange=self.exchange,
+            channel=self.channel.to_pb(),
+            symbol=self.symbol,
+        )
+
+        if self.depth is not None:
+            subscription_pb.depth = self.depth.value
+
+        if self.interval is not None:
+            subscription_pb.interval = self.interval
+
+        return subscription_pb

python/bbgo/enums/__init__.py

@@ -1,4 +1,5 @@
 from .channel_type import ChannelType
+from .depth_type import DepthType
 from .event_type import EventType
 from .order_type import OrderType
 from .side_type import SideType

python/bbgo/enums/channel_type.py

@@ -1,9 +1,32 @@
+from __future__ import annotations
+
 from enum import Enum
 
+import bbgo_pb2
+
 
 class ChannelType(Enum):
     BOOK = 'book'
     TRADE = 'trade'
     TICKER = 'ticker'
     USER = 'user'
     KLINE = 'kline'
+
+    @classmethod
+    def from_pb(cls, obj) -> ChannelType:
+        return {
+            bbgo_pb2.Channel.BOOK: cls.BOOK,
+            bbgo_pb2.Channel.TRADE: cls.TRADE,
+            bbgo_pb2.Channel.TICKER: cls.TICKER,
+            bbgo_pb2.Channel.USER: cls.USER,
+            bbgo_pb2.Channel.KLINE: cls.KLINE,
+        }[obj]
+
+    def to_pb(self) -> bbgo_pb2.Channel:
+        return {
+            'book': bbgo_pb2.Channel.BOOK,
+            'trade': bbgo_pb2.Channel.TRADE,
+            'ticker': bbgo_pb2.Channel.TICKER,
+            'user': bbgo_pb2.Channel.USER,
+            'kline': bbgo_pb2.Channel.KLINE,
+        }[self.value]

python/bbgo/enums/depth_type.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from enum import Enum
+
+
+#  string depth = 4; // depth is for book, valid values are full, medium, 1, 5 and 20
+class DepthType(Enum):
+    FULL = 'full'
+    MEDIUM = 'medium'
+    DEPTH_1 = '1'
+    DEPTH_5 = '5'
+    DEPTH_20 = '20'

python/bbgo/enums/event_type.py

@@ -1,7 +1,27 @@
+from __future__ import annotations
+
 from enum import Enum
 
+# enum Event {
+#   UNKNOWN = 0;
+#   SUBSCRIBED = 1;
+#   UNSUBSCRIBED = 2;
+#   SNAPSHOT = 3;
+#   UPDATE = 4;
+#   AUTHENTICATED = 5;
+#   ORDER_SNAPSHOT = 6;
+#   ORDER_UPDATE = 7;
+#   TRADE_SNAPSHOT = 8;
+#   TRADE_UPDATE = 9;
+#   ACCOUNT_SNAPSHOT = 10;
+#   ACCOUNT_UPDATE = 11;
+#   ERROR = 99;
+# }
+import bbgo_pb2
+
 
 class EventType(Enum):
+    UNKNOWN = 'unknown'
     ERROR = 'error'
     SUBSCRIBED = 'subscribed'
     UNSUBSCRIBED = 'unsubscribed'
@@ -14,3 +34,22 @@ class EventType(Enum):
     TRADE_UPDATE = 'trade_update'
     ACCOUNT_SNAPSHOT = 'account_snapshot'
     ACCOUNT_UPDATE = 'account_update'
+
+    @classmethod
+    def from_pb(cls, obj: bbgo_pb2.Event) -> EventType:
+        d = {
+            bbgo_pb2.Event.UNKNOWN: cls.UNKNOWN,
+            bbgo_pb2.Event.ERROR: cls.ERROR,
+            bbgo_pb2.Event.SUBSCRIBED: cls.SUBSCRIBED,
+            bbgo_pb2.Event.UNSUBSCRIBED: cls.UNSUBSCRIBED,
+            bbgo_pb2.Event.AUTHENTICATED: cls.AUTHENTICATED,
+            bbgo_pb2.Event.SNAPSHOT: cls.SNAPSHOT,
+            bbgo_pb2.Event.UPDATE: cls.UPDATE,
+            bbgo_pb2.Event.ORDER_SNAPSHOT: cls.ORDER_SNAPSHOT,
+            bbgo_pb2.Event.ORDER_UPDATE: cls.ORDER_UPDATE,
+            bbgo_pb2.Event.TRADE_SNAPSHOT: cls.TRADE_SNAPSHOT,
+            bbgo_pb2.Event.TRADE_UPDATE: cls.TRADE_UPDATE,
+            bbgo_pb2.Event.ACCOUNT_SNAPSHOT: cls.ACCOUNT_SNAPSHOT,
+            bbgo_pb2.Event.ACCOUNT_UPDATE: cls.ACCOUNT_UPDATE,
+        }
+        return d[obj]

python/bbgo/services.py

@@ -1,11 +1,13 @@
 from typing import List
-from typing import Tuple
+from typing import Tuple, Iterator
 
 import bbgo_pb2
 import bbgo_pb2_grpc
 
 from .data import ErrorMessage
+from .data import Event
 from .data import KLine
+from .data import Subscription
 
 
 class UserDataService(object):
@@ -22,10 +24,12 @@ class MarketService(object):
     def __init__(self, stub: bbgo_pb2_grpc.MarketDataServiceStub):
         self.stub = stub
 
-    def subscribe(self, subscriptions: List[bbgo_pb2.Subscription]):
-        request = bbgo_pb2.SubscribeRequest(subscriptions=subscriptions)
-        request_iter = self.stub.Subscribe(request)
-        return request_iter
+    def subscribe(self, subscriptions: List[Subscription]) -> Iterator[Event]:
+        request = bbgo_pb2.SubscribeRequest(subscriptions=[s.to_pb() for s in subscriptions])
+        response_iter = self.stub.Subscribe(request)
+
+        for response in response_iter:
+            yield Event.from_pb(response)
 
     def query_klines(self,
                      exchange: str,

python/subscribe.py

@@ -0,0 +1,30 @@
+import click
+import grpc
+from loguru import logger
+
+import bbgo_pb2_grpc
+from bbgo import MarketService
+from bbgo.data import Subscription
+from bbgo.enums import ChannelType
+from bbgo.enums import DepthType
+
+
+@click.command()
+@click.option('--host', default='127.0.0.1')
+@click.option('--port', default=50051)
+def main(host, port):
+    subscriptions = [
+        Subscription('binance', ChannelType.BOOK, symbol='BTCUSDT', depth=DepthType.FULL),
+    ]
+    address = f'{host}:{port}'
+    channel = grpc.insecure_channel(address)
+    stub = bbgo_pb2_grpc.MarketDataServiceStub(channel)
+
+    service = MarketService(stub)
+    response_iter = service.subscribe(subscriptions)
+    for response in response_iter:
+        logger.info(response)
+
+
+if __name__ == '__main__':
+    main()