Files
bitfinex-api-py/bfxapi/websockets/SubscriptionManager.py
2018-11-27 16:29:51 +00:00

103 lines
3.7 KiB
Python

import json
import asyncio
import time
from ..utils.CustomLogger import CustomLogger
from ..models import Subscription
class SubscriptionManager:
def __init__(self, bfxapi, logLevel='INFO'):
self.pending_subscriptions = {}
self.subscriptions_chanid = {}
self.subscriptions_subid = {}
self.unsubscribe_callbacks = {}
self.bfxapi = bfxapi
self.logger = CustomLogger('BfxSubscriptionManager', logLevel=logLevel)
async def subscribe(self, channel_name, symbol, timeframe=None, **kwargs):
# create a new subscription
subscription = Subscription(self.bfxapi.ws, channel_name, symbol, timeframe, **kwargs)
self.logger.info("Subscribing to channel {}".format(channel_name))
key = "{}_{}".format(channel_name, subscription.key or symbol)
self.pending_subscriptions[key] = subscription
await subscription.subscribe()
async def confirm_subscription(self, raw_ws_data):
# {"event":"subscribed","channel":"trades","chanId":1,"symbol":"tBTCUSD","pair":"BTCUSD"}
# {"event":"subscribed","channel":"candles","chanId":351,"key":"trade:1m:tBTCUSD"}
# {"event":"subscribed","channel":"book","chanId":4,"symbol":"tBTCUSD","prec":"P0","freq":"F0","len":"25","pair":"BTCUSD"}
symbol = raw_ws_data.get("symbol", None)
channel = raw_ws_data.get("channel")
chanId = raw_ws_data.get("chanId")
key = raw_ws_data.get("key", None)
get_key = "{}_{}".format(channel, key or symbol)
if chanId in self.subscriptions_chanid:
# subscription has already existed in the past
p_sub = self.subscriptions_chanid[chanId]
else:
# has just been created and is pending
p_sub = self.pending_subscriptions[get_key]
# remove from pending list
del self.pending_subscriptions[get_key]
p_sub.confirm_subscription(chanId)
# add to confirmed list
self.subscriptions_chanid[chanId] = p_sub
self.subscriptions_subid[p_sub.sub_id] = p_sub
self.bfxapi._emit('subscribed', p_sub)
async def confirm_unsubscribe(self, raw_ws_data):
chanId = raw_ws_data.get("chanId")
sub = self.subscriptions_chanid[chanId]
sub.confirm_unsubscribe()
self.bfxapi._emit('unsubscribed', sub)
# call onComplete callback if exists
if sub.sub_id in self.unsubscribe_callbacks:
await self.unsubscribe_callbacks[sub.sub_id]()
del self.unsubscribe_callbacks[sub.sub_id]
def get(self, chanId):
return self.subscriptions_chanid[chanId]
async def unsubscribe(self, chanId, onComplete=None):
sub = self.subscriptions_chanid[chanId]
if onComplete:
self.unsubscribe_callbacks[sub.sub_id] = onComplete
if sub.is_subscribed():
await self.subscriptions_chanid[chanId].unsubscribe()
async def resubscribe(self, chanId):
sub = self.subscriptions_chanid[chanId]
async def re_sub():
await sub.subscribe()
if sub.is_subscribed():
# unsubscribe first and call callback to subscribe
await self.unsubscribe(chanId, re_sub)
else:
# already unsibscribed, so just subscribe
await sub.subscribe()
def is_subscribed(self, chanId):
if chanId not in self.subscriptions_chanid:
return False
return self.subscriptions_chanid[chanId].is_subscribed()
async def unsubscribe_all(self):
task_batch = []
for chanId in self.subscriptions_chanid:
sub = self.get(chanId)
if sub.is_subscribed():
task_batch += [
asyncio.ensure_future(self.unsubscribe(chanId))
]
await asyncio.wait(*[ task_batch ])
async def resubscribe_all(self):
task_batch = []
for chanId in self.subscriptions_chanid:
task_batch += [
asyncio.ensure_future(self.resubscribe(chanId))
]
await asyncio.wait(*[ task_batch ])