47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class DmTracker:
|
||
|
|
def __init__(self):
|
||
|
|
self._direct_map: dict[str, str] = {}
|
||
|
|
self._dm_rooms: set[str] = set()
|
||
|
|
|
||
|
|
def is_direct_room(self, room_id: str) -> bool:
|
||
|
|
return room_id in self._dm_rooms
|
||
|
|
|
||
|
|
def get_dm_user(self, room_id: str) -> str | None:
|
||
|
|
return self._direct_map.get(room_id)
|
||
|
|
|
||
|
|
def add_direct_room(self, room_id: str, user_id: str) -> None:
|
||
|
|
self._dm_rooms.add(room_id)
|
||
|
|
self._direct_map[room_id] = user_id
|
||
|
|
|
||
|
|
def remove_direct_room(self, room_id: str) -> None:
|
||
|
|
self._dm_rooms.discard(room_id)
|
||
|
|
self._direct_map.pop(room_id, None)
|
||
|
|
|
||
|
|
def load_from_account_data(self, m_direct: dict[str, list[str]]) -> None:
|
||
|
|
for user_id, room_ids in m_direct.items():
|
||
|
|
for room_id in room_ids:
|
||
|
|
self._dm_rooms.add(room_id)
|
||
|
|
self._direct_map[room_id] = user_id
|
||
|
|
logger.debug("Loaded %d DM rooms from m.direct account data", len(self._dm_rooms))
|
||
|
|
|
||
|
|
def mark_candidate_dm(self, room_id: str, inviter_user_id: str) -> None:
|
||
|
|
self._dm_rooms.add(room_id)
|
||
|
|
self._direct_map[room_id] = inviter_user_id
|
||
|
|
logger.debug("Marked room %s as candidate DM with %s", room_id, inviter_user_id)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def dm_rooms(self) -> frozenset[str]:
|
||
|
|
return frozenset(self._dm_rooms)
|
||
|
|
|
||
|
|
def clear(self) -> None:
|
||
|
|
self._dm_rooms.clear()
|
||
|
|
self._direct_map.clear()
|