mirror of
https://github.com/c9s/bbgo.git
synced 2024-11-26 08:45:16 +00:00
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
import asyncio
|
|
from typing import Callable
|
|
from typing import List
|
|
|
|
import grpc
|
|
|
|
import bbgo_pb2
|
|
import bbgo_pb2_grpc
|
|
from bbgo.enums import ChannelType
|
|
from bbgo.enums import DepthType
|
|
|
|
from .data import Event
|
|
from .data import MarketDataEvent
|
|
from .data import Subscription
|
|
from .data import UserDataEvent
|
|
|
|
|
|
class Stream(object):
|
|
subscriptions: List[Subscription]
|
|
|
|
def __init__(self, host: str, port: int, user_data: bool = False):
|
|
self.host = host
|
|
self.port = port
|
|
self.user_data = user_data
|
|
|
|
self.subscriptions = []
|
|
self.event_handlers = []
|
|
|
|
def subscribe(self, exchange: str, channel: str, symbol: str, depth: str = None, interval: str = None):
|
|
subscription = Subscription(exchange=exchange, channel=ChannelType.from_str(channel), symbol=symbol)
|
|
|
|
if depth is not None:
|
|
subscription.depth = DepthType(depth)
|
|
|
|
if interval is not None:
|
|
subscription.interval = interval
|
|
|
|
self.subscriptions.append(subscription)
|
|
|
|
def add_event_handler(self, event_handler: Callable) -> None:
|
|
self.event_handlers.append(event_handler)
|
|
|
|
def fire_event_handlers(self, event: Event) -> None:
|
|
for event_handler in self.event_handlers:
|
|
event_handler(event)
|
|
|
|
@property
|
|
def address(self):
|
|
return f'{self.host}:{self.port}'
|
|
|
|
async def _subscribe(self):
|
|
async with grpc.aio.insecure_channel(self.address) as channel:
|
|
stub = bbgo_pb2_grpc.MarketDataServiceStub(channel)
|
|
|
|
request = bbgo_pb2.SubscribeRequest(subscriptions=[s.to_pb() for s in self.subscriptions])
|
|
async for response in stub.Subscribe(request):
|
|
event = MarketDataEvent.from_pb(response)
|
|
self.fire_event_handlers(event)
|
|
|
|
async def _subscribe_user_data(self):
|
|
async with grpc.aio.insecure_channel(self.address) as channel:
|
|
stub = bbgo_pb2_grpc.UserDataServiceStub(channel)
|
|
|
|
request = bbgo_pb2.Empty()
|
|
async for response in stub.SubscribeUserData(request):
|
|
event = UserDataEvent.from_pb(response)
|
|
self.fire_event_handlers(event)
|
|
|
|
def start(self):
|
|
coroutines = [self._subscribe()]
|
|
if self.user_data:
|
|
coroutines.append(self._subscribe_user_data())
|
|
|
|
group = asyncio.gather(*coroutines)
|
|
loop = asyncio.get_event_loop()
|
|
loop.run_until_complete(group)
|
|
loop.close()
|