random-tea/user.py

from __future__ import annotations
from dataclasses import dataclass, field, asdict
import typing
import time
from community import Community
from telebot.types import ChatMember

if typing.TYPE_CHECKING:
    from bot import Bot


@dataclass
class User:
    user_id: int
    chat_id: int = 0
    locale: str = ''
    start_timestamp: int = field(default_factory=lambda: int(time.time()))
    last_action_timestamp: int = field(default_factory=lambda: int(time.time()))
    # Store the chats which were checked
    # (so that we send a message for communities
    # where the user is a member of the chat but not a member of the RT community)
    checked_chats: typing.List[int] = field(default_factory=list)

    def check_chat(self, chat_id: int, bot: Bot) -> bool:
        """
        Check if the user should be notified of this chat
        """
        if chat_id in self.checked_chats:
            return False
        self.checked_chats.append(chat_id)
        chat_member_info: ChatMember = bot.get_chat_member(chat_id, self.user_id)
        return chat_member_info.is_member or chat_member_info.status in ['creator', 'administrator', 'admin', 'member']

    def dict(self) -> dict:
        data = asdict(self)
        data['last_action_timestamp'] = int(time.time())
        return data

    @classmethod
    def from_dict(cls, data: dict) -> User:
        data = getattr(data, '__dict__', data)
        data_ = {key: data[key] for key in data.keys() if
                 key in ['user_id', 'chat_id', 'locale', 'start_timestamp', 'last_action_timestamp', 'checked_chats']}
        self = cls(**data_)
        return self

    @classmethod
    def by_id(cls, user_id: int, bot) -> typing.Optional[User]:
        data = bot.get_user_from_db({'user_id': user_id})
        if data is None:
            return None
        return data

    @property
    def id(self):
        return self._id

    def is_admin(self) -> bool:
        return self.user_id in [218952152]

    def get_communities(self, bot: Bot) -> typing.Iterable[Community]:
        """
        Returns a list of community ids for which the user is a member

        Perform a MongoDB query to get the list of communities for which the user is a member
        """
        for community in bot.db.communities.find({'members': self.user_id}):
            yield Community.from_dict(community)