everyone-mention-telegram-bot/src/entity/user.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable


@dataclass
class User():
    user_id: str
    username: str
    chats: Iterable[str]

    collection: str = 'users'
    id_index: str = '_id'
    chats_index: str = 'chats'
    username_index: str = 'username'

    def is_in_chat(self, chat_id: str) -> bool:
        return chat_id in self.chats
    
    def add_to_chat(self, chat_id: str) -> None:
        self.chats.append(chat_id)

    def remove_from_chat(self, chat_id: str) -> None:
        if chat_id in self.chats:
            self.chats.remove(chat_id)

    def to_mongo_document(self) -> dict:
        return {
            self.username_index: self.username,
            self.chats_index: self.chats
        }
    
    @staticmethod
    def from_mongo_document(mongo_document: dict) -> User:
        return User(
            mongo_document[User.id_index],
            mongo_document[User.username_index],
            mongo_document[User.chats_index]
        )