craftbeerpi4-pione/cbpi/websocket.py

import logging
import weakref
from collections import defaultdict

import aiohttp
from aiohttp import web
from voluptuous import Schema

from cbpi.utils import json_dumps


class CBPiWebSocket:
    def __init__(self, cbpi) -> None:
        self.cbpi = cbpi
        self._callbacks = defaultdict(set)
        self._clients = weakref.WeakSet()
        self.logger = logging.getLogger(__name__)
        self.cbpi.app.add_routes([web.get('/ws', self.websocket_handler)])
        self.cbpi.bus.register_object(self)

        #if self.cbpi.config.static.get("ws_push_all", False):
        self.cbpi.bus.register("#", self.listen)


    async def listen(self, topic, **kwargs):
        data = dict(topic=topic, data=dict(**kwargs))
        self.logger.debug("PUSH %s " % data)
        self.send(data)


    def send(self, data, sorting=False):
        self.logger.debug("broadcast to ws clients. Data: %s" % data)
        for ws in self._clients:
            async def send_data(ws, data):
                try:
                    if sorting:
                        try:
                            data['data'].sort(key=lambda x: x.get('name').upper())
                        except:
                            pass
                    await ws.send_json(data=data, dumps=json_dumps)
                except Exception as e:
                    self.logger.error("Error with client %s: %s" % (ws, str(e)))

            self.cbpi.app.loop.create_task(send_data(ws, data))

    async def websocket_handler(self, request):
        
        

        ws = web.WebSocketResponse()
        await ws.prepare(request)
        self._clients.add(ws)
        try:
            peername = request.transport.get_extra_info('peername')
            if peername is not None:
                
                host = peername[0]
                port = peername[1]
            else:
                host, port = "Unknowen"
            self.logger.info("Client Connected - Host: %s Port: %s  - client count: %s " % (host, port, len(self._clients)))
        except Exception as e:
            pass
        
        
        try:
            await ws.send_json(data=dict(topic="connection/success"))
            async for msg in ws:
                if msg.type == aiohttp.WSMsgType.TEXT:

                    msg_obj = msg.json()
                    schema = Schema({"topic": str, "data": dict})
                    schema(msg_obj)

                    topic = msg_obj.get("topic")
                    data = msg_obj.get("data")
                    if topic == "close":
                        await ws.close()
                    else:
                        if data is not None:
                            await self.cbpi.bus.fire(topic=topic, **data)
                        else:
                            await self.cbpi.bus.fire(topic=topic)
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    self.logger.error('ws connection closed with exception %s' % ws.exception())

        except Exception as e:
            self.logger.error("%s - Received Data %s" % (str(e), msg.data))

        finally:
            self._clients.discard(ws)

        self.logger.info("Web Socket Close")

        return ws