aboutsummaryrefslogtreecommitdiffstats
path: root/ttun/pubsub.py
blob: 5d155eac49a8424cf0feb9a036c1165a8e14c28b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import asyncio
from contextlib import contextmanager
from typing import Any, Generator, Iterator


class PubSub:
    _instance = None

    def __init__(self):
        self.queues: list[asyncio.Queue] = []

    @classmethod
    def instance(cls) -> 'PubSub':
        if cls._instance == None:
            cls._instance = cls()
        return cls._instance

    @classmethod
    @contextmanager
    def subscribe(cls) -> Iterator[asyncio.Queue]:
        queue = asyncio.Queue()

        instance = cls.instance()
        instance.queues.append(queue)

        try:
            yield queue
        finally:
            instance.queues.remove(queue)
            del queue

    @classmethod
    async def publish(cls, msg: Any):
        for queue in cls.instance().queues:
            await queue.put(msg)