# Telebot imports from telebot.async_telebot import AsyncTeleBot from telebot.asyncio_storage import StateMemoryStorage, StatePickleStorage from telebot.asyncio_handler_backends import State, StatesGroup from telebot.asyncio_filters import StateFilter from telebot import types from telebot.callback_data import CallbackData, CallbackDataFilter from telebot.types import Message # Async things imports import asyncio # Other modules imports import sqlite3 import json from datetime import datetime import math import socket import os # Local imports from config import token, admins import textbook import keyboards # DB from db.base import Session, engine, Base from db.models import User, Queue, QueueUser bot = AsyncTeleBot(token, state_storage=StatePickleStorage()) class States(StatesGroup): default = State() def get_queue_stats_text(queue: Queue) -> str: s = f"Название: {queue.name}" f"Количество участников: {len(queue.users)}" return s @bot.message_handler(commands=["start"]) async def start(msg: Message): if msg.chat.type == "private": user = session.query(User).filter_by(id=msg.from_user.id).first() if not user: new_user = User( id=msg.from_user.id, name=msg.from_user.first_name, username=msg.from_user.username, ) session.add(new_user) session.commit() await bot.send_message(chat_id=msg.chat.id, text=textbook.start) await bot.set_state(user_id=msg.from_user.id, state=States.default) await bot.send_message( chat_id=msg.chat.id, text=textbook.menu, reply_markup=keyboards.menu() ) @bot.callback_query_handler(func=lambda c: c.data == "to_menu") async def to_menu_handler(call: types.CallbackQuery): await bot.set_state(user_id=call.from_user.id, state=States.default) await bot.send_message( chat_id=call.message.chat.id, text=textbook.menu, reply_markup=keyboards.menu() ) @bot.callback_query_handler(func=lambda c: c.data == "new_queue") async def new_queue_handler(call: types.CallbackQuery): user = session.query(User).filter_by(id=call.from_user.id).first() if user: if len(user.owns_queues) < 4: queue = Queue(owner_id=call.from_user.id) sessio.add(queue) session.commit() await bot.answer_callback_query( callback_query_id=call.id, text=textbook.new_queue_created.format(id=queue.id), show_alert=True, ) else: await bot.answer_callback_query( callback_query_id=call.id, text=textbook.queue_limit, show_alert=True ) await bot.answer_callback_query(callback_query_id=call.id) @bot.callback_query_handler(func=lambda c: c.data == "my_queues") async def my_queues_handler(call: types.CallbackQuery): user = session.query(User).filter_by(id=call.from_user.id).first() queues = user.owns_queues await bot.edit_message_text( chat_id=call.message.chat.id, message_id=call.message.id, text=textbook.my_queues_list.format(count=len(queues)), reply_markup=keyboards.my_queues(queues), ) await bot.answer_callback_query(callback_query_id=call.id) @bot.callback_query_handler(func=lambda c: c.data[:2] == "q:") async def edit_queue_handler(call: types.CallbackQuery): queue_id = call.data[:2] queue = session.query(Queue).filter_by(id=queue_id).first() if not queue: await bot.answer_callback_query( callback_query_id=call.id, text=textbook.error ) return None async with bot.retrieve_data( user_id=call.from_user.id, chat_id=call.message.chat.id ) as state_data: state_data["queue"] = call.data[2:] await bot.edit_message_text( chat_id=call.message.chat.id, message_id=call.message.id, text=get_queue_stats_text(queue), reply_markup=keyboards.queue_menu(), ) await bot.answer_callback_query(callback_query_id=call.id) @bot.message_handler(commands=["new_queue"]) async def nq(msg: Message): user = session.query(User).filter_by(id=msg.from_user.id).first() queue = Queue(owner_id=user.id) session.add(queue) session.commit() await bot.send_message( chat_id=msg.chat.id, text=f"Создана новая очередь: {queue.id}" ) @bot.message_handler(commands=["take_part"]) async def tp(msg: Message): try: command, queue_id = msg.text.split() except: await bot.send_message(chat_id=msg.chat.id, text="Вы забыли указать очередь") return qu = QueueUser(user_id=msg.from_user.id, queue_id=queue_id) session.add(qu) session.commit() await bot.send_message(chat_id=msg.chat.id, text="Вы приняли участие в очереди!") @bot.message_handler(commands=["queue"]) async def q(msg: Message): try: command, queue_id = msg.text.split() except: await bot.send_message(chat_id=msg.chat.id, text="Вы забыли указать очередь") return queue = session.query(Queue).filter_by(id=queue_id).first() if queue: users = [q.user.name for q in queue.users] await bot.send_message( chat_id=msg.chat.id, text=f"Пользователи этой очереди: {', '.join(users)}" ) else: await bot.send_message(chat_id=msg.chat.id, text="Очередь не найдена!") async def main(): a = asyncio.create_task(bot.polling(non_stop=True)) await a if __name__ == "__main__": print("Bot started", flush=True) Base.metadata.create_all(engine) session = Session() bot.add_custom_filter(StateFilter(bot)) bot.enable_saving_states(filename="./.state-save/states.pkl") asyncio.run(main())