bitfinex-api-py/bfxapi/websockets/generic_websocket.py

"""
Module used as a interfeace to describe a generick websocket client
"""

import asyncio
import websockets
import socket
import json
import time
from threading import Thread

from pyee import EventEmitter
from ..utils.custom_logger import CustomLogger

# websocket exceptions
from websockets.exceptions import ConnectionClosed

class AuthError(Exception):
    """
    Thrown whenever there is a problem with the authentication packet
    """
    pass

def is_json(myjson):
    try:
        json_object = json.loads(myjson)
    except ValueError as e:
        return False
    return True

class Socket():
    def __init__(self, sId):
        self.ws = None
        self.isConnected = False
        self.isAuthenticated = False
        self.id = sId

    def set_connected(self):
        self.isConnected = True

    def set_disconnected(self):
        self.isConnected = False

    def set_authenticated(self):
        self.isAuthenticated = True

    def set_websocket(self, ws):
        self.ws = ws

def _start_event_worker():
    async def event_sleep_process():
        """
        sleeping process for event emitter to schedule on
        """
        while True:
            await asyncio.sleep(0)
    def start_loop(loop):
        asyncio.set_event_loop(loop)
        loop.run_until_complete(event_sleep_process())
    event_loop = asyncio.new_event_loop()
    worker = Thread(target=start_loop, args=(event_loop,))
    worker.start()
    return event_loop

class GenericWebsocket:
    """
    Websocket object used to contain the base functionality of a websocket.
    Inlcudes an event emitter and a standard websocket client.
    """

    def __init__(self, host, logLevel='INFO', loop=None, max_retries=5,
                 create_event_emitter=_start_event_worker):
        self.host = host
        self.logger = CustomLogger('BfxWebsocket', logLevel=logLevel)
        self.loop = loop or asyncio.get_event_loop()
        # overide 'error' event to stop it raising an exception
        # self.events.on('error', self.on_error)
        self.ws = None
        self.max_retries = max_retries
        self.attempt_retry = True
        self.sockets = {}
        # start seperate process for the even emitter
        eventLoop = create_event_emitter()
        self.events = EventEmitter(scheduler=asyncio.ensure_future, loop=eventLoop)

    def run(self):
        """
        Starte the websocket connection. This functions spawns the initial socket
        thread and connection.
        """
        self._start_new_socket()

    def get_task_executable(self):
        """
        Get the run indefinitely asyncio task
        """
        return self._run_socket()

    def _start_new_socket(self, socketId=None):
        if not socketId:
            socketId = len(self.sockets)
        def start_loop(loop):
            asyncio.set_event_loop(loop)
            loop.run_until_complete(self._run_socket())
        worker_loop = asyncio.new_event_loop()
        worker = Thread(target=start_loop, args=(worker_loop,))
        worker.start()
        return socketId

    def _wait_for_socket(self, socket_id):
        """
        Block until the given socket connection is open
        """
        while True:
            socket = self.sockets.get(socket_id, False)
            if socket:
                if socket.isConnected and socket.ws:
                    return
            time.sleep(0.01)

    async def _connect(self, socket):
        async with websockets.connect(self.host) as websocket:
            self.sockets[socket.id].set_websocket(websocket)
            self.sockets[socket.id].set_connected()
            self.logger.info("Wesocket connected to {}".format(self.host))
            while True:
                await asyncio.sleep(0)
                message = await websocket.recv()
                await self.on_message(socket.id, message)

    def get_socket(self, socketId):
        return self.sockets[socketId]

    def get_authenticated_socket(self):
        for socketId in self.sockets:
            if self.sockets[socketId].isAuthenticated:
                return self.sockets[socketId]
        return None

    async def _run_socket(self):
        retries = 0
        sId =  len(self.sockets)
        s = Socket(sId)
        self.sockets[sId] = s
        while retries < self.max_retries and self.attempt_retry:
            try:
                await self._connect(s)
                retries = 0
            except (ConnectionClosed, socket.error) as e:
                self.sockets[sId].set_disconnected()
                self._emit('disconnected')
                if (not self.attempt_retry):
                    return
                self.logger.error(str(e))
                retries += 1
                # wait 5 seconds befor retrying
                self.logger.info("Waiting 5 seconds before retrying...")
                await asyncio.sleep(5)
                self.logger.info("Reconnect attempt {}/{}".format(retries, self.max_retries))
        self.logger.info("Unable to connect to websocket.")
        self._emit('stopped')

    def remove_all_listeners(self, event):
        """
        Remove all listeners from event emitter
        """
        self.events.remove_all_listeners(event)

    def on(self, event, func=None):
        """
        Add a new event to the event emitter
        """
        if not func:
            return self.events.on(event)
        self.events.on(event, func)

    def once(self, event, func=None):
        """
        Add a new event to only fire once to the event
        emitter
        """
        if not func:
            return self.events.once(event)
        self.events.once(event, func)

    def _emit(self, event, *args, **kwargs):
        self.events.emit(event, *args, **kwargs)

    async def on_error(self, error):
        """
        On websocket error print and fire event
        """
        self.logger.error(error)

    async def on_close(self):
        """
        On websocket close print and fire event. This is used by the data server.
        """
        self.logger.info("Websocket closed.")
        self.attempt_retry = False
        await self.ws.close()
        self._emit('done')

    async def on_open(self):
        """
        On websocket open
        """
        pass

    async def on_message(self, message):
        """
        On websocket message
        """
        pass