chuchube-py/channel.py

from typing import Dict

from websockets.asyncio.server import ServerConnection
from websockets.exceptions import ConnectionClosed

class Subscriber:
    player_enabled = False
    ws: ServerConnection

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


class Channel:
    subscribers: Dict[ServerConnection, Subscriber]

    def __init__(self):
        self.subscribers = dict()

    def subscribe(self, ws: ServerConnection):
        if ws not in self.subscribers:
            self.subscribers[ws] = (Subscriber(ws))

    def unsubscribe(self, ws: ServerConnection):
        if ws in self.subscribers:
            self.subscribers.pop(ws)

    def get_player_enabled_subscribers(self):
        return filter(lambda s: s.player_enabled, self.subscribers.values())

    async def send(self, message):
        for sub in list(self.subscribers.values()):
            try:
                await sub.ws.send(message)
            except ConnectionClosed:
                self.unsubscribe(sub.ws)
            # else:
            #     print("closed ws still in channel")
            #     self.unsubscribe(sub.ws)