Source code for divera.socket

import asyncio

import socketio


class SocketConnection:
    def __init__(
            self,
            jwt: [str, callable],
            ucr: [str, callable],
            config: dict = None,
    ):
        self.jwt = jwt
        self.ucr = ucr

        self.client = socketio.AsyncClient(
            **(config or {}),
        )

    async def authenticate(
            self,
    ):
        if asyncio.iscoroutinefunction(self.jwt):
            jwt = await self.jwt('ws')
        elif callable(self.jwt):
            jwt = self.jwt()
        else:
            jwt = self.jwt

        if asyncio.iscoroutinefunction(self.ucr):
            ucr = await self.ucr()
        elif callable(self.ucr):
            ucr = self.ucr()
        else:
            ucr = self.ucr

        await self.client.emit(
            event='authenticate',
            data={
                'jwt': jwt,
                'ucr': ucr,
            },
        )

    def add_callback(
            self,
            callback: callable,
            event: str = '*',
            namespace: str = None,
    ):
        self.client.on(
            event=event,
            handler=callback,
            namespace=namespace,
        )

    async def setup(self):
        @self.client.event
        async def connect():
            await self.authenticate()

    async def connect(
            self,
            url: str,
    ):
        await self.setup()
        await self.client.connect(
            url=url,
        )