mautrix-python/mautrix/appservice/state_store/memory.py

70 lines
2.2 KiB
Python

# Copyright (c) 2022 Tulir Asokan
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
from typing import Dict, Optional, Tuple
from abc import ABC
import time
from mautrix.client.state_store import StateStore as ClientStateStore
from mautrix.types import EventID, RoomID, UserID
class ASStateStore(ClientStateStore, ABC):
_presence: Dict[UserID, str]
_read: Dict[Tuple[RoomID, UserID], EventID]
_registered: Dict[UserID, bool]
def __init__(self) -> None:
self._registered = {}
# Non-persistent storage
self._presence = {}
self._read = {}
async def is_registered(self, user_id: UserID) -> bool:
"""
Check if a given user is registered.
This should always return ``True`` for double puppets, because they're always registered
beforehand and shouldn't be attempted to register by the bridge.
Args:
user_id: The user ID to check.
Returns:
``True`` if the user is registered, ``False`` otherwise.
"""
if not user_id:
raise ValueError("user_id is empty")
return self._registered.get(user_id, False)
async def registered(self, user_id: UserID) -> None:
"""
Mark the given user ID as registered.
Args:
user_id: The user ID to mark as registered.
"""
if not user_id:
raise ValueError("user_id is empty")
self._registered[user_id] = True
def set_presence(self, user_id: UserID, presence: str) -> None:
self._presence[user_id] = presence
def has_presence(self, user_id: UserID, presence: str) -> bool:
try:
return self._presence[user_id] == presence
except KeyError:
return False
def set_read(self, room_id: RoomID, user_id: UserID, event_id: EventID) -> None:
self._read[(room_id, user_id)] = event_id
def get_read(self, room_id: RoomID, user_id: UserID) -> Optional[EventID]:
try:
return self._read[(room_id, user_id)]
except KeyError:
return None