websockets: adds event driven

bfxapi/websockets/GenericWebsocket.py

@@ -2,6 +2,7 @@ import asyncio
 import websockets
 import json
 
+from pyee import EventEmitter
 from ..utils.CustomLogger import CustomLogger
 
 class AuthError(Exception): pass
@@ -14,23 +15,17 @@ def is_json(myjson):
   return True
 
 class GenericWebsocket(object):
-  def __init__(self, host, symbol='tBTCUSD', onCandleHook=None, onTradeHook=None, onCompleteHook=None):
-    if not onCandleHook:
-      raise KeyError("Expected `onCandleHook` in parameters.")
-    if not onTradeHook:
-      raise KeyError("Expected `onTradeHook` in parameters.")
-    if not onCompleteHook:
-      raise KeyError("Expected `onCompleteHook` in parameters.")
-    self.onCandleHook = onCandleHook
-    self.onTradeHook = onTradeHook
-    self.onCompleteHook = onCompleteHook
+
+  def __init__(self, host, symbol='tBTCUSD', onCandleHook=None, onTradeHook=None, 
+      logLevel='ERROR'):
     self.symbol = symbol
     self.host = host
     self.awaiting_request = False
-
-    self.logger = CustomLogger('HFWebSocket', logLevel='INFO')
+    self.onCandleHook = onCandleHook
+    self.onTradeHook = onTradeHook
+    self.logger = CustomLogger('HFWebSocket', logLevel=logLevel)
     self.loop = asyncio.get_event_loop()
-    # self.events = EventEmitter()
+    self.events = EventEmitter(scheduler=asyncio.ensure_future, loop=self.loop)
 
   def run(self):
     self.loop.run_until_complete(self._main(self.host))
@@ -43,13 +38,25 @@ class GenericWebsocket(object):
         message = await websocket.recv()
         await self.on_message(message)
 
+  def on(self, event, func=None):
+    if not func:
+      return self.events.on(event)
+    self.events.on(event, func)
+
+  def _emit(self, event, *args, **kwargs):
+    self.events.emit(event, *args, **kwargs)
+
+  def once(self, event, func):
+    self.events.once(event, func)
+
   async def on_error(self, error):
     self.logger.error(error)
+    self.events.emit('error', error)
 
   async def on_close(self):
     self.logger.info("Websocket closed.")
     await self.ws.close()
-    self.onCompleteHook()
+    self._emit('done')
 
   async def on_open(self):
     pass