import telebot import pymongo from functools import wraps from user import User from templates import TemplateProvider import emoji_data_python as edp from threading import Thread import utils import typing import traceback import time def locale_from_ietf(ietf_locale: str) -> str: if not ietf_locale: return 'en' ietf_locale = ietf_locale.split('-')[0] return 'ru' if ietf_locale in ['ru', 'be', 'uk'] else 'en' class Bot(telebot.TeleBot): def __init__(self, config: dict): self.name = config['name'] super().__init__(config['token']) self.me: telebot.types.User = self.get_me() self.db = getattr(pymongo.MongoClient(config['mongo']), self.name) self.db.users.create_index([('user_id', pymongo.ASCENDING)], unique=True) self.config = config self.templates = TemplateProvider() for template_file in config.get('template_files', []): self.templates.load_file(template_file) self.main_keyboard = None if config.get('main_keyboard') is not None: kb = config.get('main_keyboard') self.main_keyboard = {key: utils.create_keyboard(kb[key]) for key in kb.keys()} def render_template(self, template_name: str, locale: str = 'ru', **kwargs) -> \ typing.Tuple[str, typing.Union[telebot.types.ReplyKeyboardMarkup, telebot.types.InlineKeyboardMarkup]]: text: str = self.templates.render_template(template_name, locale, bot=self, **kwargs) text, keyboard = self.templates.separate_text_and_keyboards(text) keyboard = utils.create_keyboard(keyboard) \ if keyboard is not None \ else self.main_keyboard.get(locale, self.main_keyboard.get('ru')) return text, keyboard def render_message_for_user(self, user: User, template_name: str, **kwargs) -> \ typing.Tuple[str, typing.Union[telebot.types.ReplyKeyboardMarkup, telebot.types.InlineKeyboardMarkup]]: text: str = self.templates.render_template(template_name, user.locale, user=user, **kwargs) text, keyboard = self.templates.separate_text_and_keyboards(text) keyboard = utils.create_keyboard(keyboard) if keyboard is not None \ else self.main_keyboard.get(user.locale, self.main_keyboard.get('ru')) return text, keyboard def edit_message_with_template(self, query: telebot.types.CallbackQuery, template: str, **kwargs): msg = query.message text, keyboard = self.render_message_for_user(self.get_user_by_msg(msg), template, tg_user=msg.from_user, **kwargs) self.edit_message_text(text, msg.chat.id, msg.id, reply_markup=keyboard, parse_mode='HTML') self.answer_callback_query(query) def reply_with_template(self, msg: typing.Union[telebot.types.Message, telebot.types.CallbackQuery], template: str, as_reply: bool = False, locale_overwrite: typing.Optional[str] = None, **kwargs): if isinstance(msg, telebot.types.CallbackQuery): msg = msg.message text, keyboard = self.render_message_for_user(self.get_user_by_msg(msg), template, tg_user=msg.from_user, locale_overwrite=locale_overwrite, **kwargs) if as_reply: self.reply_to(msg, text, reply_markup=keyboard) else: self.send_message(msg.chat.id, text, reply_markup=keyboard, parse_mode='HTML') def get_user_from_db(self, request: dict) -> typing.Optional[User]: data = self.db.users.find_one(request) if data is None: return None return User.from_dict(data) def get_user_by_msg(self, msg: typing.Union[telebot.types.Message, telebot.types.InlineQuery]) -> User: tg_user: telebot.types.User = msg.from_user if tg_user.id == self.me.id and msg.chat.id > 0: user = User.by_id(msg.chat.id, self) or User.by_id(msg.from_user.id, self) else: user = User.by_id(msg.from_user.id, self) if user is not None: if tg_user.language_code: user.locale = locale_from_ietf(tg_user.language_code) return user else: user = User(msg.from_user.id, msg.from_user.id) user.locale = locale_from_ietf(tg_user.language_code) self.db.users.insert_one(user.dict()) return user def save(self, user: User): self.db.users.replace_one({'user_id': user.user_id}, user.dict()) def handle_commands(self, cmds: typing.List[str]): cmds = [edp.replace_colons(cmd) for cmd in cmds] def wrapper(func): @wraps(func) def func_wrapped(msg): user = self.get_user_by_msg(msg) args: str = msg.text for cmd in cmds: if args.startswith(cmd): if len(args.split(cmd)) <= 1: args = '' break args = cmd.join(args.split(cmd)[1:]).strip() break func(msg, user, args) if cmds[0].startswith('/'): self.message_handler(commands=[cmd[1:] for cmd in cmds])(func_wrapped) else: self.message_handler(func=lambda msg: any(msg.text.startswith(cmd) for cmd in cmds))(func_wrapped) return wrapper def handle_callback(self, name: str): def wrapper(func): @wraps(func) def func_wrapped(query): user = self.get_user_by_msg(query) args = query.data.strip()[len(name) + 1:] try: return func(query, user, args.strip()) except Exception: print(f'An exception occured: {traceback.format_exc()}') self.callback_query_handler( lambda query: query.data.strip().startswith(name + '_') or query.data.strip() == name)(func_wrapped) return func_wrapped return wrapper def inline_handler_(self, filter_func, **kwargs): def wrapper(func): @wraps(func) def handler(*args, **kwargs_2): def func_2(): try: func(*args, **kwargs_2) except Exception as e: print('An exception occured:', e) Thread(target=func_2).start() self.inline_handler(filter_func, **kwargs)(func) return handler return wrapper def iter_users(self) -> typing.Iterable[User]: for user_data in self.db.users.find(): yield User.from_dict(user_data) def register_next_step_handler(self, message, callback, *args, **kwargs): if isinstance(message, telebot.types.CallbackQuery): message = message.message def new_callback(new_message, *new_args, **new_kwargs): try: if new_message.text in ['๐Ÿ›‘ Cancel', '๐Ÿ›‘ ะžั‚ะผะตะฝะธั‚ัŒ']: # self.clear_step_handler_by_chat_id(new_message.chat.id) self.reply_with_template(new_message, 'canceled') return callback(new_message, *new_args, **new_kwargs) except: print(traceback.format_exc()) super().register_next_step_handler(message, new_callback, *args, **kwargs) def send_all_copy(self, message: telebot.types.Message): num = 0 failures = 0 for user in self.iter_users(): try: self.copy_message(user.user_id, message.chat.id, message.id) num += 1 except: print(traceback.format_exc()) failures += 1 print(f'Sent to {num} users. Failures: {failures}') def user_growth(self) -> typing.List[int]: timestamp = next(iter(self.db.users.find({}).sort('start_timestamp'))) data = list() while timestamp < time.time(): data.append(self.db.users.find({'start_timestamp': {'$lt': timestamp}}).count_documents()) timestamp += 24 * 3600 return data def stats(self) -> typing.Dict[str, int]: total_users = self.db.users.count_documents({}) new_users_today = self.db.users.count_documents( {'start_timestamp': {'$gt': int(time.time() - 24 * 3600)}} ) new_users_month = self.db.users.count_documents( {'start_timestamp': {'$gt': int(time.time() - 24 * 3600 * 30)}} ) active_users_today = self.db.users.count_documents( {'last_action_timestamp': {'$gt': int(time.time() - 24 * 3600)}} ) active_users_month = self.db.users.count_documents( {'last_action_timestamp': {'$gt': int(time.time() - 24 * 3600 * 30)}} ) return { 'Total users': total_users, 'New users in the last 24 hours': new_users_today, 'New users last 30 days': new_users_month, 'Active users in the last 24 hours': active_users_today, 'Active users in the last 30 days': active_users_month }