import logging from dataclasses import dataclass from typing import Union from bson import ObjectId from libbot import sync from pyrogram.types import User from classes.pyroclient import PyroClient from modules.database import col_groups logger = logging.getLogger(__name__) @dataclass class PyroGroup: """Dataclass of DB entry of a group""" __slots__ = ("_id", "id", "locale", "locale_auto") _id: ObjectId id: int locale: Union[str, None] locale_auto: bool ban_failed: bool @classmethod async def create_if_not_exists( cls, id: int, locale: Union[str, None] = sync.config_get("locale", "defaults", "group"), locale_auto: bool = sync.config_get("locale_auto", "defaults", "group"), ban_failed: bool = sync.config_get("ban_failed", "defaults", "group"), ): db_entry = await col_groups.find_one( { "id": id, } ) if db_entry is None: inserted = await col_groups.insert_one( { "id": id, "locale": locale, "locale_auto": locale_auto, "ban_failed": ban_failed, } ) db_entry = await col_groups.find_one({"_id": inserted.inserted_id}) if db_entry is None: raise RuntimeError("Could not find inserted group entry.") return cls(**db_entry) async def set_locale(self, locale: Union[str, None]) -> None: logger.debug("Locale of group %s has been set to %s", self.id, locale) await col_groups.update_one({"_id": self._id}, {"$set": {"locale": locale}}) async def set_locale_auto(self, enabled: bool) -> None: logger.debug( "Automatic locale selection of group %s has been set to %s", self.id, enabled, ) await col_groups.update_one( {"_id": self._id}, {"$set": {"locale_auto": enabled}} ) async def set_ban_failed(self, enabled: bool) -> None: logger.debug( "Banning users that failed the captcha in group %s has been set to %s", self.id, enabled, ) await col_groups.update_one( {"_id": self._id}, {"$set": {"ban_failed": enabled}} ) # Group settings # User locale def select_locale( self, app: PyroClient, user: User, ignore_auto: bool = False ) -> str: if not ignore_auto and self.locale_auto is True: if user.language_code is not None: return user.language_code return self.locale if self.locale is not None else app.default_locale return self.locale if self.locale is not None else app.default_locale