Merge pull request #526 from narumiruna/python/fix-enums

python: fix enums

python/bbgo/data/event.py

@@ -32,8 +32,8 @@ class UserDataEvent(Event):
         return cls(
             session=obj.session,
             exchange=obj.exchange,
-            channel_type=ChannelType.from_pb(obj.channel),
-            event_type=EventType.from_pb(obj.event),
+            channel_type=ChannelType(obj.channel),
+            event_type=EventType(obj.event),
             balances=[Balance.from_pb(balance) for balance in obj.balances],
             trades=[Trade.from_pb(trade) for trade in obj.trades],
             orders=[Order.from_pb(order) for order in obj.orders],
@@ -50,13 +50,13 @@ class MarketDataEvent(Event):
 
     @classmethod
     def from_pb(cls, obj: bbgo_pb2.SubscribeResponse) -> MarketDataEvent:
-        channel_type = ChannelType.from_pb(obj.channel)
+        channel_type = ChannelType(obj.channel)
 
         event = cls(
             exchange=obj.exchange,
             symbol=obj.symbol,
             channel_type=channel_type,
-            event_type=EventType.from_pb(obj.event),
+            event_type=EventType(obj.event),
         )
 
         if channel_type == ChannelType.BOOK:

python/bbgo/data/order.py

@@ -31,8 +31,8 @@ class Order:
             exchange=obj.exchange,
             symbol=obj.symbol,
             order_id=obj.id,
-            side=SideType.from_pb(obj.side),
-            order_type=OrderType.from_pb(obj.order_type),
+            side=SideType(obj.side),
+            order_type=OrderType(obj.order_type),
             price=float(obj.price),
             stop_price=float(obj.stop_price),
             status=obj.status,

python/bbgo/data/subscription.py

@@ -19,7 +19,7 @@ class Subscription:
     def to_pb(self) -> bbgo_pb2.Subscription:
         subscription_pb = bbgo_pb2.Subscription(
             exchange=self.exchange,
-            channel=self.channel.to_pb(),
+            channel=self.channel.value,
             symbol=self.symbol,
         )
 

python/bbgo/enums/channel_type.py

@@ -2,28 +2,15 @@ from __future__ import annotations
 
 from enum import Enum
 
-import bbgo_pb2
-
 
 class ChannelType(Enum):
-    BOOK = 'book'
-    TRADE = 'trade'
-    TICKER = 'ticker'
-    KLINE = 'kline'
+    BOOK = 0
+    TRADE = 1
+    TICKER = 2
+    KLINE = 3
+    BALANCE = 4
+    ORDER = 5
 
-    @staticmethod
-    def from_pb(obj) -> ChannelType:
-        return {
-            bbgo_pb2.Channel.BOOK: ChannelType.BOOK,
-            bbgo_pb2.Channel.TRADE: ChannelType.TRADE,
-            bbgo_pb2.Channel.TICKER: ChannelType.TICKER,
-            bbgo_pb2.Channel.KLINE: ChannelType.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,
-            'kline': bbgo_pb2.Channel.KLINE,
-        }[self.value]
+    @classmethod
+    def from_str(cls, s: str) -> ChannelType:
+        return {t.name.lower(): t for t in cls}[s.lower()]

python/bbgo/enums/depth_type.py

@@ -10,3 +10,7 @@ class DepthType(Enum):
     DEPTH_1 = '1'
     DEPTH_5 = '5'
     DEPTH_20 = '20'
+
+    @classmethod
+    def from_str(cls, s: str) -> DepthType:
+        return {t.name.lower(): t for t in cls}[s.lower()]

python/bbgo/enums/event_type.py

@@ -2,27 +2,16 @@ from __future__ import annotations
 
 from enum import Enum
 
-import bbgo_pb2
-
 
 class EventType(Enum):
-    UNKNOWN = 'unknown'
-    ERROR = 'error'
-    SUBSCRIBED = 'subscribed'
-    UNSUBSCRIBED = 'unsubscribed'
-    AUTHENTICATED = 'authenticated'
-    SNAPSHOT = 'snapshot'
-    UPDATE = 'update'
+    UNKNOWN = 0
+    SUBSCRIBED = 1
+    UNSUBSCRIBED = 2
+    SNAPSHOT = 3
+    UPDATE = 4
+    AUTHENTICATED = 5
+    ERROR = 99
 
-    @staticmethod
-    def from_pb(obj: bbgo_pb2.Event) -> EventType:
-        d = {
-            bbgo_pb2.Event.UNKNOWN: EventType.UNKNOWN,
-            bbgo_pb2.Event.ERROR: EventType.ERROR,
-            bbgo_pb2.Event.SUBSCRIBED: EventType.SUBSCRIBED,
-            bbgo_pb2.Event.UNSUBSCRIBED: EventType.UNSUBSCRIBED,
-            bbgo_pb2.Event.AUTHENTICATED: EventType.AUTHENTICATED,
-            bbgo_pb2.Event.SNAPSHOT: EventType.SNAPSHOT,
-            bbgo_pb2.Event.UPDATE: EventType.UPDATE,
-        }
-        return d[obj]
+    @classmethod
+    def from_str(cls, s: str) -> EventType:
+        return {t.name.lower(): t for t in cls}[s.lower()]

python/bbgo/enums/order_type.py

@@ -2,24 +2,15 @@ from __future__ import annotations
 
 from enum import Enum
 
-import bbgo_pb2
-
 
 class OrderType(Enum):
-    MARKET = 'market'
-    LIMIT = 'limit'
-    STOP_MARKET = 'stop_market'
-    STOP_LIMIT = 'stop_limit'
-    POST_ONLY = 'post_only'
-    IOC_LIMIT = 'ioc_limit'
+    MARKET = 0
+    LIMIT = 1
+    STOP_MARKET = 2
+    STOP_LIMIT = 3
+    POST_ONLY = 4
+    IOC_LIMIT = 5
 
-    @staticmethod
-    def from_pb(obj: bbgo_pb2.OrderType) -> OrderType:
-        return {
-            bbgo_pb2.OrderType.MARKET: OrderType.MARKET,
-            bbgo_pb2.OrderType.LIMIT: OrderType.LIMIT,
-            bbgo_pb2.OrderType.STOP_MARKET: OrderType.STOP_MARKET,
-            bbgo_pb2.OrderType.STOP_LIMIT: OrderType.STOP_LIMIT,
-            bbgo_pb2.OrderType.POST_ONLY: OrderType.POST_ONLY,
-            bbgo_pb2.OrderType.IOC_LIMIT: OrderType.IOC_LIMIT,
-        }[obj]
+    @classmethod
+    def from_str(cls, s: str) -> OrderType:
+        return {t.name.lower(): t for t in cls}[s.lower()]

python/bbgo/enums/side_type.py

@@ -2,16 +2,11 @@ from __future__ import annotations
 
 from enum import Enum
 
-import bbgo_pb2
-
 
 class SideType(Enum):
-    BUY = 'buy'
-    SELL = 'sell'
+    BUY = 0
+    SELL = 1
 
-    @staticmethod
-    def from_pb(obj: bbgo_pb2.Side) -> SideType:
-        return {
-            bbgo_pb2.Side.BUY: SideType.BUY,
-            bbgo_pb2.Side.SELL: SideType.SELL,
-        }[obj]
+    @classmethod
+    def from_str(cls, s: str) -> SideType:
+        return {t.name.lower(): t for t in cls}[s.lower()]

python/bbgo/stream.py

@@ -27,7 +27,7 @@ class Stream(object):
         self.event_handlers = []
 
     def subscribe(self, exchange: str, channel: str, symbol: str, depth: str = None, interval: str = None):
-        subscription = Subscription(exchange=exchange, channel=ChannelType(channel), symbol=symbol)
+        subscription = Subscription(exchange=exchange, channel=ChannelType.from_str(channel), symbol=symbol)
 
         if depth is not None:
             subscription.depth = DepthType(depth)

python/examples/subscribe_user_data.py

@@ -0,0 +1,24 @@
+import grpc
+from loguru import logger
+
+import bbgo_pb2
+import bbgo_pb2_grpc
+from bbgo.data import UserDataEvent
+
+
+def main():
+    host = '127.0.0.1'
+    port = 50051
+    address = f'{host}:{port}'
+    channel = grpc.insecure_channel(address)
+    stub = bbgo_pb2_grpc.UserDataServiceStub(channel)
+
+    request = bbgo_pb2.UserDataRequest(session='max')
+    response_iter = stub.Subscribe(request)
+    for response in response_iter:
+        event = UserDataEvent.from_pb(response)
+        logger.info(event)
+
+
+if __name__ == '__main__':
+    main()

python/tests/test_enums.py

@@ -0,0 +1,15 @@
+from bbgo.enums import ChannelType
+
+
+def test_channel_type_from_str():
+    m = {
+        'book': ChannelType.BOOK,
+        'trade': ChannelType.TRADE,
+        'ticker': ChannelType.TICKER,
+        'kline': ChannelType.KLINE,
+        'balance': ChannelType.BALANCE,
+        'order': ChannelType.ORDER,
+    }
+
+    for k, v in m.items():
+        assert ChannelType.from_str(k) == v