49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
import asyncio
|
|
import random
|
|
import json
|
|
|
|
_behaviors_by_species = {
|
|
"dog": ["wagging_tail", "stretching", "showing_belly", "barking", "panting"],
|
|
"cat": ["blep", "stretching", "showing_belly", "purring", "kneading"],
|
|
"bird": ["head_tilt", "stretching", "chirping", "wing_flap"],
|
|
}
|
|
_subscribers: dict[int, list[asyncio.Queue]] = {}
|
|
_worker_tasks: dict[int, asyncio.Task] = {}
|
|
|
|
|
|
async def _behavior_worker(pet_id: int, species: str):
|
|
behaviors = _behaviors_by_species.get(species, ["wagging_tail"])
|
|
while True:
|
|
action = random.choice(behaviors)
|
|
data = json.dumps({"action": action})
|
|
queues = _subscribers.get(pet_id, [])
|
|
for q in queues:
|
|
await q.put(data)
|
|
await asyncio.sleep(random.uniform(1, 5))
|
|
|
|
|
|
def _ensure_worker(pet_id: int, species: str):
|
|
if pet_id not in _worker_tasks or _worker_tasks[pet_id].done():
|
|
_subscribers.setdefault(pet_id, [])
|
|
_worker_tasks[pet_id] = asyncio.create_task(
|
|
_behavior_worker(pet_id, species)
|
|
)
|
|
|
|
|
|
async def subscribe(pet_id: int, species: str) -> asyncio.Queue:
|
|
q: asyncio.Queue = asyncio.Queue()
|
|
_subscribers.setdefault(pet_id, []).append(q)
|
|
_ensure_worker(pet_id, species)
|
|
return q
|
|
|
|
|
|
def unsubscribe(pet_id: int, q: asyncio.Queue):
|
|
queues = _subscribers.get(pet_id, [])
|
|
if q in queues:
|
|
queues.remove(q)
|
|
if not queues:
|
|
task = _worker_tasks.pop(pet_id, None)
|
|
if task and not task.done():
|
|
task.cancel()
|
|
_subscribers.pop(pet_id, None)
|