You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
23 lines
641 B
Python
23 lines
641 B
Python
2 years ago
|
import asyncio
|
||
|
|
from typing import AsyncGenerator
|
||
|
|
|
||
|
|
from lws.models import Message
|
||
|
|
|
||
|
|
class Broker:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.connections = set()
|
||
|
|
|
||
|
|
async def publish(self, message: str, store=True) -> None:
|
||
|
|
for connection in self.connections:
|
||
|
|
Message(message=message).save()
|
||
|
|
await connection.put(message)
|
||
|
|
|
||
|
|
async def subscribe(self) -> AsyncGenerator[str, None]:
|
||
|
|
connection = asyncio.Queue()
|
||
|
|
self.connections.add(connection)
|
||
|
|
try:
|
||
|
|
while True:
|
||
|
|
yield await connection.get()
|
||
|
|
finally:
|
||
|
|
self.connections.remove(connection)
|