From de2c905144da8a09663d4286ef7690ea39ff75f0 Mon Sep 17 00:00:00 2001 From: krau Date: Tue, 10 Oct 2023 21:52:14 +0800 Subject: [PATCH] =?UTF-8?q?=E9=87=8D=E6=9E=84=E9=A1=B9=E7=9B=AE=E7=BB=93?= =?UTF-8?q?=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Dockerfile | 2 +- docker-compose.yml | 2 +- kmua/__init__.py | 0 bot.py => kmua/__main__.py | 15 +- kmua/callbacks/__init__.py | 0 {src => kmua}/callbacks/chatdata.py | 11 +- {src => kmua}/callbacks/chatinfo.py | 7 +- {src => kmua}/callbacks/chatmember.py | 21 +-- {src => kmua}/callbacks/friendship.py | 8 +- {src => kmua}/callbacks/help.py | 10 +- {src => kmua}/callbacks/jobs.py | 9 +- {src => kmua}/callbacks/keyword_reply.py | 14 +- {src => kmua}/callbacks/manage.py | 66 +++---- {src => kmua}/callbacks/others.py | 8 +- {src => kmua}/callbacks/quote.py | 71 +++----- {src => kmua}/callbacks/remake.py | 17 +- {src => kmua}/callbacks/slash.py | 6 +- {src => kmua}/callbacks/start.py | 25 ++- {src => kmua}/callbacks/sticker.py | 2 +- {src => kmua}/callbacks/title.py | 6 +- {src => kmua}/callbacks/userdata.py | 66 +++---- {src => kmua}/callbacks/waifu.py | 163 +++++++----------- kmua/common/__init__.py | 7 + {src => kmua}/common/bot.py | 15 +- {src => kmua}/common/chat.py | 6 +- {src => kmua}/common/message.py | 9 +- {src => kmua}/common/quote.py | 8 +- {src => kmua}/common/user.py | 27 ++- {src => kmua}/common/utils.py | 4 +- {src => kmua}/common/waifu.py | 22 +-- src/config/__init__.py => kmua/config.py | 4 +- kmua/dao/__init__.py | 7 + src/dao/db.py => kmua/dao/_db.py | 13 +- {src => kmua}/dao/association.py | 26 +-- {src => kmua}/dao/chat.py | 52 +++--- kmua/dao/chat_service.py | 40 +++++ {src => kmua}/dao/quote.py | 12 +- {src => kmua}/dao/user.py | 18 +- .../user.py => kmua/dao/user_service.py | 27 ++- {src/service => kmua/dao}/waifu.py | 122 ++++++------- {src => kmua}/filters.py | 0 {src => kmua}/handlers.py | 155 ++++++++--------- {src => kmua}/logger.py | 2 +- {src/models => kmua}/models.py | 11 +- {resource => kmua/resource}/TsukuA.ttc | Bin {resource => kmua/resource}/base.png | Bin .../resource}/word_dicts/weni.json | 0 pyproject.toml | 4 +- src/dao/__init__.py | 13 -- src/service/chat.py | 35 ---- 50 files changed, 533 insertions(+), 635 deletions(-) create mode 100644 kmua/__init__.py rename bot.py => kmua/__main__.py (89%) create mode 100644 kmua/callbacks/__init__.py rename {src => kmua}/callbacks/chatdata.py (71%) rename {src => kmua}/callbacks/chatinfo.py (95%) rename {src => kmua}/callbacks/chatmember.py (86%) rename {src => kmua}/callbacks/friendship.py (87%) rename {src => kmua}/callbacks/help.py (87%) rename {src => kmua}/callbacks/jobs.py (87%) rename {src => kmua}/callbacks/keyword_reply.py (89%) rename {src => kmua}/callbacks/manage.py (80%) rename {src => kmua}/callbacks/others.py (86%) rename {src => kmua}/callbacks/quote.py (87%) rename {src => kmua}/callbacks/remake.py (64%) rename {src => kmua}/callbacks/slash.py (96%) rename {src => kmua}/callbacks/start.py (84%) rename {src => kmua}/callbacks/sticker.py (97%) rename {src => kmua}/callbacks/title.py (97%) rename {src => kmua}/callbacks/userdata.py (84%) rename {src => kmua}/callbacks/waifu.py (78%) create mode 100644 kmua/common/__init__.py rename {src => kmua}/common/bot.py (70%) rename {src => kmua}/common/chat.py (82%) rename {src => kmua}/common/message.py (87%) rename {src => kmua}/common/quote.py (88%) rename {src => kmua}/common/user.py (92%) rename {src => kmua}/common/utils.py (97%) rename {src => kmua}/common/waifu.py (92%) rename src/config/__init__.py => kmua/config.py (74%) create mode 100644 kmua/dao/__init__.py rename src/dao/db.py => kmua/dao/_db.py (53%) rename {src => kmua}/dao/association.py (76%) rename {src => kmua}/dao/chat.py (69%) create mode 100644 kmua/dao/chat_service.py rename {src => kmua}/dao/quote.py (85%) rename {src => kmua}/dao/user.py (83%) rename src/service/user.py => kmua/dao/user_service.py (50%) rename {src/service => kmua/dao}/waifu.py (61%) rename {src => kmua}/filters.py (100%) rename {src => kmua}/handlers.py (50%) rename {src => kmua}/logger.py (88%) rename {src/models => kmua}/models.py (93%) rename {resource => kmua/resource}/TsukuA.ttc (100%) rename {resource => kmua/resource}/base.png (100%) rename {resource => kmua/resource}/word_dicts/weni.json (100%) delete mode 100644 src/dao/__init__.py delete mode 100644 src/service/chat.py diff --git a/Dockerfile b/Dockerfile index d4f9fdb..a5bd14e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,4 +2,4 @@ FROM python:3.11.6-slim-bookworm COPY . /kmua WORKDIR /kmua RUN apt-get update && apt-get install graphviz -y && pip install -r requirements.txt -ENTRYPOINT [ "python","/kmua/bot.py" ] +ENTRYPOINT [ "python","-m","kmua" ] diff --git a/docker-compose.yml b/docker-compose.yml index a4510d7..310c536 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -2,7 +2,7 @@ version: "3" services: kmua: image: ghcr.io/krau/kmua-bot:v2 - container_name: kmua-v2-main + container_name: kmua-main restart: always volumes: - ./data:/kmua/data # 不要修改此挂载 diff --git a/kmua/__init__.py b/kmua/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bot.py b/kmua/__main__.py similarity index 89% rename from bot.py rename to kmua/__main__.py index 8c2a5b7..709d0ed 100644 --- a/bot.py +++ b/kmua/__main__.py @@ -1,5 +1,4 @@ import datetime -from src.dao import db import pytz from telegram.constants import UpdateType @@ -9,11 +8,11 @@ ApplicationBuilder, Defaults, ) - -from src.callbacks.jobs import refresh_waifu_data -from src.config import settings -from src.handlers import handlers, on_error -from src.logger import logger +from kmua.logger import logger +from kmua.config import settings +import kmua.dao._db as db +from kmua.handlers import handlers, on_error +from kmua.callbacks.jobs import refresh_waifu_data async def init_data(app: Application): @@ -34,7 +33,6 @@ async def init_data(app: Application): ("id", "获取聊天ID"), ] ) - # dao.init_db() logger.success("started bot") @@ -80,5 +78,4 @@ def run(): ) -if __name__ == "__main__": - run() +run() diff --git a/kmua/callbacks/__init__.py b/kmua/callbacks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/callbacks/chatdata.py b/kmua/callbacks/chatdata.py similarity index 71% rename from src/callbacks/chatdata.py rename to kmua/callbacks/chatdata.py index b970938..505a04d 100644 --- a/src/callbacks/chatdata.py +++ b/kmua/callbacks/chatdata.py @@ -2,15 +2,16 @@ Update, ) from telegram.ext import ContextTypes -from ..common.chat import get_chat_info -from ..logger import logger -from ..service import chat as chat_service + +import kmua.common as common +from kmua.logger import logger +import kmua.dao as dao async def chat_data_manage(update: Update, context: ContextTypes.DEFAULT_TYPE): chat = update.effective_chat logger.info(f"chat_data_manage: {chat.title}") - text = get_chat_info(chat) + text = common.get_chat_info(chat) await chat.send_message(text=text) # TODO: manage chat data @@ -19,4 +20,4 @@ async def chat_title_update(update: Update, context: ContextTypes.DEFAULT_TYPE): chat = update.effective_chat logger.info(f"update_chat_title: {chat.title}") title = chat.title - chat_service.update_chat_title(chat=chat, title=title) + dao.update_chat_title(chat=chat, title=title) diff --git a/src/callbacks/chatinfo.py b/kmua/callbacks/chatinfo.py similarity index 95% rename from src/callbacks/chatinfo.py rename to kmua/callbacks/chatinfo.py index 0828503..0a00d35 100644 --- a/src/callbacks/chatinfo.py +++ b/kmua/callbacks/chatinfo.py @@ -1,9 +1,8 @@ -from telegram import ( - Update, -) -from ..logger import logger +from telegram import Update from telegram.ext import ContextTypes +from kmua.logger import logger + async def getid(update: Update, context: ContextTypes.DEFAULT_TYPE): logger.info( diff --git a/src/callbacks/chatmember.py b/kmua/callbacks/chatmember.py similarity index 86% rename from src/callbacks/chatmember.py rename to kmua/callbacks/chatmember.py index 2ccf490..8e4054d 100644 --- a/src/callbacks/chatmember.py +++ b/kmua/callbacks/chatmember.py @@ -3,12 +3,9 @@ from telegram import Chat, ChatMember, ChatMemberUpdated, Update from telegram.ext import ContextTypes -from ..dao.association import delete_association_in_chat -from ..dao.chat import add_chat -from ..dao import db -from ..logger import logger -from ..service.chat import delete_chat_data_and_quotes -from ..common.user import verify_user_can_manage_bot_in_chat +import kmua.common as common +import kmua.dao as dao +from kmua.logger import logger def extract_status_change(chat_member_update: ChatMemberUpdated): @@ -60,7 +57,7 @@ async def track_chats(update: Update, context: ContextTypes.DEFAULT_TYPE) -> Non elif was_member and not is_member: logger.debug(f"{cause_name} 将bot移出群组 {chat.title}") - delete_chat_data_and_quotes(chat) + dao.delete_chat_data_and_quotes(chat) elif not was_member and is_member: logger.debug(f"{cause_name} 将bot添加到频道 {chat.title}") @@ -77,7 +74,7 @@ async def track_chats(update: Update, context: ContextTypes.DEFAULT_TYPE) -> Non async def on_member_left(update: Update, context: ContextTypes.DEFAULT_TYPE): left_user = update.effective_message.left_chat_member logger.debug(f"{left_user.full_name} 退出了群聊 {update.effective_chat.title}") - delete_association_in_chat(update.effective_chat, left_user) + dao.delete_association_in_chat(update.effective_chat, left_user) async def on_member_join(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -86,7 +83,7 @@ async def on_member_join(update: Update, context: ContextTypes.DEFAULT_TYPE): logger.debug(f"{joined_user} 加入了群聊 {chat.title} ") if joined_user.is_bot: return - db_chat = add_chat(chat) + db_chat = dao.add_chat(chat) greet = db_chat.greet if greet: greet = greet.format( @@ -105,13 +102,11 @@ async def set_greet(update: Update, context: ContextTypes.DEFAULT_TYPE): chat = update.effective_chat logger.info(f"[{chat.title}]({user.name}) ") message = update.effective_message - if not await verify_user_can_manage_bot_in_chat(user, chat, update, context): + if not await common.verify_user_can_manage_bot_in_chat(user, chat, update, context): await message.reply_text("你没有权限哦") return greet = " ".join(context.args) - db_chat = add_chat(chat) - db_chat.greet = greet - db.commit() + dao.update_chat_greet(chat, greet) if greet: await message.reply_text("设置成功") else: diff --git a/src/callbacks/friendship.py b/kmua/callbacks/friendship.py similarity index 87% rename from src/callbacks/friendship.py rename to kmua/callbacks/friendship.py index 7480b20..cb857c1 100644 --- a/src/callbacks/friendship.py +++ b/kmua/callbacks/friendship.py @@ -3,9 +3,9 @@ from telegram import Update from telegram.ext import ContextTypes -from ..common.utils import ohayo_word, oyasumi_word -from ..logger import logger +from kmua.logger import logger from .jobs import send_message +import kmua.common as common async def ohayo(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -18,7 +18,7 @@ async def ohayo(update: Update, context: ContextTypes.DEFAULT_TYPE): when=random.randint(22, 25) * random.randint(3500, 3600), data={ "chat_id": update.effective_user.id, - "text": random.choice(ohayo_word), + "text": random.choice(common.ohayo_word), }, name=f"ohayo_{update.effective_user.id}", ) @@ -34,7 +34,7 @@ async def oyasumi(update: Update, context: ContextTypes.DEFAULT_TYPE): when=random.randint(22, 25) * random.randint(3500, 3600), data={ "chat_id": update.effective_user.id, - "text": random.choice(oyasumi_word), + "text": random.choice(common.oyasumi_word), }, name=f"oyasumi_{update.effective_user.id}", ) diff --git a/src/callbacks/help.py b/kmua/callbacks/help.py similarity index 87% rename from src/callbacks/help.py rename to kmua/callbacks/help.py index 22295e8..e8f78da 100644 --- a/src/callbacks/help.py +++ b/kmua/callbacks/help.py @@ -5,8 +5,8 @@ ) from telegram.ext import ContextTypes -from ..logger import logger -from ..common.message import message_recorder +import kmua.common as common +from kmua.logger import logger async def help(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -48,10 +48,10 @@ async def help(update: Update, context: ContextTypes.DEFAULT_TYPE): ] ] ) - sent_message = await context.bot.send_message( + await context.bot.send_message( chat_id=update.effective_chat.id, text=help_text, reply_markup=help_markup, ) - message_recorder(update, context) - logger.info(f"Bot: {sent_message.text}") + common.message_recorder(update, context) + logger.info("Bot: ") diff --git a/src/callbacks/jobs.py b/kmua/callbacks/jobs.py similarity index 87% rename from src/callbacks/jobs.py rename to kmua/callbacks/jobs.py index 89abcaf..e7dcf6e 100644 --- a/src/callbacks/jobs.py +++ b/kmua/callbacks/jobs.py @@ -2,9 +2,8 @@ from telegram.ext import ContextTypes -from ..dao.chat import get_all_chats -from ..logger import logger -from ..service.waifu import refresh_all_waifu_data +from kmua.logger import logger +import kmua.dao as dao from .waifu import send_waifu_graph @@ -13,7 +12,7 @@ async def refresh_waifu_data(context: ContextTypes.DEFAULT_TYPE): try: context.bot_data["refeshing_waifu_data"] = True await asyncio.gather( - *(send_waifu_graph(chat, context) for chat in get_all_chats()) + *(send_waifu_graph(chat, context) for chat in dao.get_all_chats()) ) except Exception as err: logger.error( @@ -22,7 +21,7 @@ async def refresh_waifu_data(context: ContextTypes.DEFAULT_TYPE): raise err finally: await asyncio.sleep(3) - await refresh_all_waifu_data() + await dao.refresh_all_waifu_data() logger.success("数据已刷新: waifu_data") context.bot_data["refeshing_waifu_data"] = False diff --git a/src/callbacks/keyword_reply.py b/kmua/callbacks/keyword_reply.py similarity index 89% rename from src/callbacks/keyword_reply.py rename to kmua/callbacks/keyword_reply.py index 68b73f3..69cccd8 100644 --- a/src/callbacks/keyword_reply.py +++ b/kmua/callbacks/keyword_reply.py @@ -6,12 +6,14 @@ from telegram.constants import ChatAction from telegram.ext import ContextTypes from zhconv import convert + import asyncio -from ..common.utils import word_dict -from ..logger import logger -from ..common.message import message_recorder -from .friendship import ohayo, oyasumi + +import kmua.common as common +from kmua.logger import logger from telegram.error import BadRequest +from .friendship import ohayo, oyasumi + async def keyword_reply(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -28,7 +30,7 @@ async def keyword_reply(update: Update, context: ContextTypes.DEFAULT_TYPE): ) sent_message = None all_resplist = [] - for keyword, resplist in word_dict.items(): + for keyword, resplist in common.word_dict.items(): if keyword in message_text: all_resplist.extend(resplist) if keyword == "早": @@ -42,7 +44,7 @@ async def keyword_reply(update: Update, context: ContextTypes.DEFAULT_TYPE): allow_sending_without_reply=False, ) logger.info("Bot: " + sent_message.text) - message_recorder(update, context) + common.message_recorder(update, context) if sent_message: await asyncio.sleep(30) source_message = sent_message.reply_to_message diff --git a/src/callbacks/manage.py b/kmua/callbacks/manage.py similarity index 80% rename from src/callbacks/manage.py rename to kmua/callbacks/manage.py index 7a0fea7..899d56b 100644 --- a/src/callbacks/manage.py +++ b/kmua/callbacks/manage.py @@ -1,27 +1,15 @@ -import asyncio - -from telegram import ( - InlineKeyboardButton, - InlineKeyboardMarkup, - Update, -) +from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup from telegram.ext import ContextTypes +from kmua.logger import logger +import kmua.common as common -from ..common.bot import get_bot_status -from ..common.user import ( - download_big_avatar, - fake_users_id, - verify_user_can_manage_bot, - verify_user_can_manage_bot_in_chat, -) -from ..dao import db -from ..dao.association import get_association_in_chat_by_user -from ..dao.user import get_user_by_id -from ..logger import logger -from ..service.user import check_user_in_chat +import asyncio +import kmua.dao as dao from .chatdata import chat_data_manage + from .jobs import refresh_waifu_data + _manage_markup = InlineKeyboardMarkup( [[InlineKeyboardButton("Refresh bot info", callback_data="bot_data_refresh")]] ) @@ -33,7 +21,7 @@ async def manage(update: Update, context: ContextTypes.DEFAULT_TYPE): await chat_data_manage(update, context) return # TODO: manage bot in private chat - if not verify_user_can_manage_bot(update.effective_user): + if not common.verify_user_can_manage_bot(update.effective_user): return if context.bot_data.get("lock_manage_bot", False): await chat.send_message("Locked") @@ -56,15 +44,15 @@ async def bot_data_refresh(update: Update, context: ContextTypes.DEFAULT_TYPE): query = update.callback_query try: await query.answer("正在刷新 bot 数据...") - db_bot_user = get_user_by_id(context.bot.id) - avatar = await download_big_avatar(context.bot.id, context) + db_bot_user = dao.get_user_by_id(context.bot.id) + avatar = await common.download_big_avatar(context.bot.id, context) db_bot_user.avatar_big_blob = avatar sent_message = await chat.send_photo( photo=avatar, caption="此消息用于获取 bot 头像缓存 id", ) db_bot_user.avatar_big_id = sent_message.photo[-1].file_id - db.commit() + dao.commit() await sent_message.delete() await query.edit_message_text("刷新完成") finally: @@ -80,12 +68,12 @@ async def set_bot_admin_in_chat(update: Update, context: ContextTypes.DEFAULT_TY chat = update.effective_chat logger.info(f"[{chat.title}]({user.name}) ") message = update.effective_message - if not await verify_user_can_manage_bot_in_chat(user, chat, update, context): + if not await common.verify_user_can_manage_bot_in_chat(user, chat, update, context): await update.message.reply_text("你没有权限哦") return to_promote_user_id = None if replied_message := message.reply_to_message: - if replied_message.from_user.id in fake_users_id: + if replied_message.from_user.id in common.fake_users_id: await message.reply_text("暂不支持设置该类用户为bot管理员") return to_promote_user_id = replied_message.from_user.id @@ -94,24 +82,24 @@ async def set_bot_admin_in_chat(update: Update, context: ContextTypes.DEFAULT_TY try: to_promote_user_id = int(to_promote_user_id) except ValueError: - await message.reply_text("请输入正确的用户ID, 或者回复一名消息") + await message.reply_text("请输入正确的用户ID, 或者回复一条消息") return else: - await message.reply_text("请输入正确的用户ID, 或者回复一名消息") + await message.reply_text("请输入正确的用户ID, 或者回复一条消息") return - db_user = get_user_by_id(to_promote_user_id) + db_user = dao.get_user_by_id(to_promote_user_id) if not db_user: await message.reply_text("该用户不存在") return - if not check_user_in_chat(db_user, chat): + if not dao.check_user_in_chat(db_user, chat): await message.reply_text("该用户不在本群") return if not db_user.is_real_user: await message.reply_text("暂不支持设置该类用户为bot管理员") return - association = get_association_in_chat_by_user(chat, db_user) + association = dao.get_association_in_chat_by_user(chat, db_user) association.is_bot_admin = not association.is_bot_admin - db.commit() + dao.commit() await message.reply_text( f"已将{db_user.full_name}在本群的bot管理权限设置为{association.is_bot_admin}" ) @@ -120,7 +108,7 @@ async def set_bot_admin_in_chat(update: Update, context: ContextTypes.DEFAULT_TY async def set_bot_admin_globally(update: Update, context: ContextTypes.DEFAULT_TYPE): user = update.effective_user logger.info(f"{user.name} ") - if not verify_user_can_manage_bot(user): + if not common.verify_user_can_manage_bot(user): return message = update.effective_message if not context.args: @@ -132,12 +120,12 @@ async def set_bot_admin_globally(update: Update, context: ContextTypes.DEFAULT_T except ValueError: await message.reply_text("请输入正确的用户ID") return - db_user = get_user_by_id(to_promote_user_id) + db_user = dao.get_user_by_id(to_promote_user_id) if not db_user: await message.reply_text("该用户不存在") return db_user.is_bot_global_admin = not db_user.is_bot_global_admin - db.commit() + dao.commit() await message.reply_text( f"已将{db_user.full_name}的bot全局管理权限设置为{db_user.is_bot_global_admin}" ) @@ -146,7 +134,7 @@ async def set_bot_admin_globally(update: Update, context: ContextTypes.DEFAULT_T async def leave_chat(update: Update, context: ContextTypes.DEFAULT_TYPE): user = update.effective_user logger.info(f"[{user.name}] ") - if not verify_user_can_manage_bot(user): + if not common.verify_user_can_manage_bot(user): return chat = update.effective_chat if chat.type == chat.GROUP or chat.type == chat.SUPERGROUP: @@ -174,7 +162,7 @@ async def refresh_waifu_data_manually( update: Update, context: ContextTypes.DEFAULT_TYPE ): user = update.effective_user - if not verify_user_can_manage_bot(user): + if not common.verify_user_can_manage_bot(user): return await update.effective_message.reply_text("3s 后刷新 waifu_data") context.job_queue.run_once(refresh_waifu_data, 3) @@ -194,11 +182,13 @@ async def status(update: Update, context: ContextTypes.DEFAULT_TYPE): try: if not query: await update.effective_message.reply_text( - get_bot_status(), reply_markup=_status_markup + common.get_bot_status(), reply_markup=_status_markup ) return - await query.edit_message_text(get_bot_status(), reply_markup=_status_markup) + await query.edit_message_text( + common.get_bot_status(), reply_markup=_status_markup + ) finally: await asyncio.sleep(1) context.user_data["lock_status"] = False diff --git a/src/callbacks/others.py b/kmua/callbacks/others.py similarity index 86% rename from src/callbacks/others.py rename to kmua/callbacks/others.py index a565265..3b940f1 100644 --- a/src/callbacks/others.py +++ b/kmua/callbacks/others.py @@ -1,9 +1,9 @@ from telegram import Update from telegram.ext import ContextTypes -from ..config import settings -from ..logger import logger -from ..service.chat import update_chat_id +from kmua.logger import logger +from kmua.config import settings +import kmua.dao as dao async def chat_migration(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -18,7 +18,7 @@ async def chat_migration(update: Update, context: ContextTypes.DEFAULT_TYPE): new_chat_id = message.migrate_to_chat_id if old_chat_id is None or new_chat_id is None: return - update_chat_id(old_chat_id, new_chat_id) + dao.update_chat_id(old_chat_id, new_chat_id) async def error_notice_control(update: Update, context: ContextTypes.DEFAULT_TYPE): diff --git a/src/callbacks/quote.py b/kmua/callbacks/quote.py similarity index 87% rename from src/callbacks/quote.py rename to kmua/callbacks/quote.py index 380faa7..e68e559 100644 --- a/src/callbacks/quote.py +++ b/kmua/callbacks/quote.py @@ -14,30 +14,9 @@ from telegram.ext import ContextTypes from telegram.helpers import escape_markdown -from ..common.message import ( - get_message_common_link, - message_recorder, - parse_message_link, -) -from ..common.quote import generate_quote_img -from ..common.user import ( - get_big_avatar_bytes, - verify_user_can_manage_bot_in_chat, - verify_user_is_chat_admin, -) -from ..common.utils import random_unit -from ..dao.chat import ( - get_chat_quote_probability, - get_chat_quotes, - get_chat_quotes_count, - get_chat_quotes_page, - update_chat_quote_probability, -) -from ..dao.quote import add_quote, delete_quote, delete_quote_by_link, get_quote_by_link -from ..dao.user import ( - add_user, -) -from ..logger import logger +import kmua.common as common +import kmua.dao as dao +from kmua.logger import logger from .jobs import delete_message @@ -51,7 +30,7 @@ async def quote(update: Update, context: ContextTypes.DEFAULT_TYPE): message = update.effective_message logger.info(f"[{chat.title}]({user.name})" + f" {message.text}") - message_recorder(update, context) + common.message_recorder(update, context) if chat.type == ChatType.PRIVATE: return @@ -68,9 +47,9 @@ async def quote(update: Update, context: ContextTypes.DEFAULT_TYPE): quote_user = forward_from_user qer_user = message.sender_chat or user - add_user(quote_user) - quote_message_link = get_message_common_link(quote_message) - if get_quote_by_link(quote_message_link): + dao.add_user(quote_user) + quote_message_link = common.get_message_common_link(quote_message) + if dao.get_quote_by_link(quote_message_link): sent_message = await message.reply_markdown_v2( "这条消息已经在语录中了哦" "\n_This message will be deleted in 10s_" ) @@ -128,11 +107,11 @@ async def _generate_and_sned_quote_img( return None if len(quote_message.text) > 200: return None - avatar = await get_big_avatar_bytes(quote_user.id, context) + avatar = await common.get_big_avatar_bytes(quote_user.id, context) if not avatar: return None await update.effective_chat.send_action(ChatAction.UPLOAD_PHOTO) - quote_img = await generate_quote_img( + quote_img = await common.generate_quote_img( avatar=avatar, text=quote_message.text, name=quote_user.title if isinstance(quote_user, Chat) else quote_user.name, @@ -157,7 +136,7 @@ async def _save_quote_data( :param quote_user: 语录消息的发送者 :param quote_img: 语录图片的 file_id """ - add_quote( + dao.add_quote( chat=update.effective_chat, user=quote_user, qer=qer, @@ -175,9 +154,9 @@ async def set_quote_probability(update: Update, context: ContextTypes.DEFAULT_TY chat = update.effective_chat message = update.effective_message logger.info(f"[{chat.title}]({user.name})" + f" {message.text}") - message_recorder(update, context) + common.message_recorder(update, context) - if not await verify_user_is_chat_admin(user, chat, context): + if not await common.verify_user_is_chat_admin(user, chat, context): sent_message = await message.reply_text("你没有权限哦") logger.info(f"Bot: {sent_message.text}") return @@ -192,7 +171,7 @@ async def set_quote_probability(update: Update, context: ContextTypes.DEFAULT_TY sent_message = await message.reply_text(except_text) logger.info(f"Bot: {sent_message.text}") return - update_chat_quote_probability(chat, probability) + dao.update_chat_quote_probability(chat, probability) sent_message = await message.reply_text( text=f"将本聊天的 random quote 的概率设为{probability}啦", ) @@ -232,16 +211,16 @@ async def random_quote(update: Update, context: ContextTypes.DEFAULT_TYPE): f"[{chat.title}({chat.id})]<{user.name}>" + (f" {message.text}" if message.text else "<非文本消息>") ) - message_recorder(update, context) + common.message_recorder(update, context) - pb = get_chat_quote_probability(chat) - flag = random_unit(pb) + pb = dao.get_chat_quote_probability(chat) + flag = common.random_unit(pb) if message.text is not None: if message.text.startswith("/qrand"): flag = True if not flag: return - quotes = get_chat_quotes(chat) + quotes = dao.get_chat_quotes(chat) if not quotes: return quote = random.choice(quotes) @@ -254,7 +233,7 @@ async def random_quote(update: Update, context: ContextTypes.DEFAULT_TYPE): logger.info(f"Bot: {sent_message.text}") except Exception as e: logger.warning(f"{e.__class__.__name__}: {e}") - delete_quote(quote) + dao.delete_quote(quote) await _unpin_messsage(message_id, chat.id, context) @@ -268,11 +247,11 @@ async def delete_quote_in_chat(update: Update, context: ContextTypes.DEFAULT_TYP message = update.effective_message if not update.callback_query: logger.info(f"[{chat.title}]({user.name})" + f" {message.text}") - message_recorder(update, context) + common.message_recorder(update, context) query_data = "" if update.callback_query: query_data = update.callback_query.data - user_can_manage_bot = await verify_user_can_manage_bot_in_chat( + user_can_manage_bot = await common.verify_user_can_manage_bot_in_chat( user, chat, update, context ) if (not message.reply_to_message and user_can_manage_bot) or ( @@ -299,13 +278,13 @@ async def delete_quote_in_chat(update: Update, context: ContextTypes.DEFAULT_TYP message_id = quote_message.id chat_id = chat.id - link = get_message_common_link(quote_message) + link = common.get_message_common_link(quote_message) if "delete_quote_in_chat" in query_data: link = query_data.split(" ")[1] - chat_id, message_id = parse_message_link(link) + chat_id, message_id = common.parse_message_link(link) await _unpin_messsage(message_id, chat_id, context) - if delete_quote_by_link(link): + if dao.delete_quote_by_link(link): if not update.callback_query: sent_message = await quote_message.reply_text("已删除该语录") logger.info(f"Bot: {sent_message.text}") @@ -336,10 +315,10 @@ async def _chat_quote_manage(update: Update, context: ContextTypes.DEFAULT_TYPE) int(update.callback_query.data.split(" ")[-1]) if update.callback_query else 1 ) page_size = 5 - quotes = get_chat_quotes_page( + quotes = dao.get_chat_quotes_page( chat=update.effective_chat, page=page, page_size=page_size ) - quotes_count = get_chat_quotes_count(chat) + quotes_count = dao.get_chat_quotes_count(chat) max_page = ceil(quotes_count / page_size) if quotes_count == 0 and not update.callback_query: sent_message = await message.reply_text("本群还没有语录哦") diff --git a/src/callbacks/remake.py b/kmua/callbacks/remake.py similarity index 64% rename from src/callbacks/remake.py rename to kmua/callbacks/remake.py index b95bd06..9912e14 100644 --- a/src/callbacks/remake.py +++ b/kmua/callbacks/remake.py @@ -4,12 +4,11 @@ from telegram import Update from telegram.ext import ContextTypes -from ..logger import logger -from ..common.message import message_recorder -from ..common.utils import random_unit +from kmua.logger import logger +import kmua.common as common -country = [ +_country = [ "中国", "日本", "韩国", @@ -23,8 +22,8 @@ "阿根廷", "印度", ] -role = ["男孩子", "女孩子", "薯条", "xyn", "猫猫", "狗狗", "鼠鼠"] -birthplace = ["首都", "省会", "直辖市", "市区", "县城", "自治区", "农村", "大学"] +_role = ["男孩子", "女孩子", "薯条", "xyn", "猫猫", "狗狗", "鼠鼠"] +_birthplace = ["首都", "省会", "直辖市", "市区", "县城", "自治区", "农村", "大学"] async def remake(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -35,14 +34,14 @@ async def remake(update: Update, context: ContextTypes.DEFAULT_TYPE): f"[{update.effective_chat.title}]({update.effective_user.name})" + f" {update.effective_message.text}" ) - if random_unit(0.114): + if common.random_unit(0.114): await update.effective_message.reply_text(text="重开失败!您没能出生!") return - text = f"重开成功\! 您出生在*{random.choice(country)}*的*{random.choice(birthplace)}*\! 是*{random.choice(role)}*\!" # noqa: E501 + text = f"重开成功\! 您出生在*{random.choice(_country)}*的*{random.choice(_birthplace)}*\! 是*{random.choice(_role)}*\!" # noqa: E501 sent_message = await update.effective_message.reply_text( text=text, parse_mode="MarkdownV2" ) - message_recorder(update, context) + common.message_recorder(update, context) logger.info(f"Bot: {sent_message.text}") await asyncio.sleep(30) context.user_data["remake_cd"] = False diff --git a/src/callbacks/slash.py b/kmua/callbacks/slash.py similarity index 96% rename from src/callbacks/slash.py rename to kmua/callbacks/slash.py index 266b1c8..f35343d 100644 --- a/src/callbacks/slash.py +++ b/kmua/callbacks/slash.py @@ -3,8 +3,8 @@ ) from telegram.ext import ContextTypes from telegram.helpers import escape_markdown -from ..logger import logger -from ..common.message import message_recorder +from kmua.logger import logger +import kmua.common as common import re @@ -94,4 +94,4 @@ async def slash(update: Update, context: ContextTypes.DEFAULT_TYPE): parse_mode="MarkdownV2", disable_web_page_preview=True, ) - message_recorder(update, context) + common.message_recorder(update, context) diff --git a/src/callbacks/start.py b/kmua/callbacks/start.py similarity index 84% rename from src/callbacks/start.py rename to kmua/callbacks/start.py index 67d9a3f..2446983 100644 --- a/src/callbacks/start.py +++ b/kmua/callbacks/start.py @@ -1,15 +1,10 @@ from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update from telegram.ext import ContextTypes -from ..logger import logger -from ..common.message import message_recorder +from kmua.logger import logger +import kmua.common as common -from ..dao import db -from ..dao.user import ( - add_user, - get_user_by_id, -) -from ..common.user import download_big_avatar +import kmua.dao as dao _start_bot_markup = InlineKeyboardMarkup( [ @@ -31,7 +26,7 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE): logger.info(f"[{update.effective_user.name}] ") - message_recorder(update, context) + common.message_recorder(update, context) if update.effective_chat.type != "private": if update.effective_message.text == "/start": # 如果是群聊，且没有艾特，直接返回 @@ -50,11 +45,13 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE): text="Nya~", reply_markup=_start_bot_markup, ) - db_bot_user = get_user_by_id(context.bot.id) + db_bot_user = dao.get_user_by_id(context.bot.id) if not db_bot_user: - db_bot_user = add_user((await context.bot.get_me())) - db_bot_user.avatar_big_blob = await download_big_avatar(context.bot.id, context) - db.commit() + db_bot_user = dao.add_user((await context.bot.get_me())) + db_bot_user.avatar_big_blob = await common.download_big_avatar( + context.bot.id, context + ) + dao.commit() photo = db_bot_user.avatar_big_id or db_bot_user.avatar_big_blob sent_message = await update.effective_message.reply_photo( photo=photo, @@ -62,7 +59,7 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE): reply_markup=_start_bot_markup, ) db_bot_user.avatar_big_id = sent_message.photo[-1].file_id - db.commit() + dao.commit() async def _start_in_group(update: Update, context: ContextTypes.DEFAULT_TYPE): diff --git a/src/callbacks/sticker.py b/kmua/callbacks/sticker.py similarity index 97% rename from src/callbacks/sticker.py rename to kmua/callbacks/sticker.py index 723d898..d5c6cf8 100644 --- a/src/callbacks/sticker.py +++ b/kmua/callbacks/sticker.py @@ -1,7 +1,7 @@ from telegram import Update from telegram.ext import ContextTypes from telegram.constants import ChatAction -from ..logger import logger +from kmua.logger import logger async def sticker2img(update: Update, context: ContextTypes.DEFAULT_TYPE): diff --git a/src/callbacks/title.py b/kmua/callbacks/title.py similarity index 97% rename from src/callbacks/title.py rename to kmua/callbacks/title.py index b7e7cb2..28303cc 100644 --- a/src/callbacks/title.py +++ b/kmua/callbacks/title.py @@ -4,7 +4,7 @@ from telegram.error import BadRequest from telegram.ext import ContextTypes from telegram.helpers import escape_markdown -from ..common.user import verify_user_can_manage_bot_in_chat +import kmua.common as common from ..logger import logger @@ -142,7 +142,7 @@ async def set_title_permissions(update: Update, context: ContextTypes.DEFAULT_TY chat = update.effective_chat logger.info(f"[{chat.title}]({user.name}) ") message = update.effective_message - if not await verify_user_can_manage_bot_in_chat(user, chat, update, context): + if not await common.verify_user_can_manage_bot_in_chat(user, chat, update, context): await message.reply_text("你没有权限哦") return title_permissions = context.chat_data.get("title_permissions", {}) @@ -170,7 +170,7 @@ async def set_title_permissions_callback( user = update.effective_user chat = update.effective_chat logger.info(f"[{chat.title}]({user.name}) ") - if not await verify_user_can_manage_bot_in_chat(user, chat, update, context): + if not await common.verify_user_can_manage_bot_in_chat(user, chat, update, context): return permission = update.callback_query.data.split(" ")[1] title_permissions = context.chat_data.get("title_permissions", {}) diff --git a/src/callbacks/userdata.py b/kmua/callbacks/userdata.py similarity index 84% rename from src/callbacks/userdata.py rename to kmua/callbacks/userdata.py index ef51e13..da37451 100644 --- a/src/callbacks/userdata.py +++ b/kmua/callbacks/userdata.py @@ -4,25 +4,11 @@ from telegram.ext import ContextTypes from telegram.helpers import escape_markdown -from ..common.user import ( - get_user_info, - download_big_avatar, - download_small_avatar, -) -from ..common.utils import back_home_markup -from ..common.waifu import get_user_waifu_info -from ..dao import db -from ..dao.quote import delete_quote_by_link -from ..dao.user import ( - get_user_by_id, - get_user_quotes_count, - get_user_quotes_page, -) -from ..logger import logger -from ..service.waifu import ( - refresh_user_all_waifu, -) from .jobs import reset_user_cd +from kmua.logger import logger +import kmua.dao as dao +import kmua.common as common + _user_data_manage_markup = InlineKeyboardMarkup( [ @@ -38,8 +24,8 @@ async def user_data_manage(update: Update, context: ContextTypes.DEFAULT_TYPE): user = update.effective_user query = update.callback_query logger.info(f"({user.name}) ") - db_user = get_user_by_id(user.id) - info = get_user_info(user) + db_user = dao.get_user_by_id(user.id) + info = common.get_user_info(user) if db_user.avatar_big_id: await query.edit_message_media( media=InputMediaPhoto( @@ -85,8 +71,8 @@ async def user_data_refresh(update: Update, context: ContextTypes.DEFAULT_TYPE): await query.answer("刷新中...") username = user.username full_name = user.full_name - avatar_big_blog = await download_big_avatar(user.id, context) - avatar_small_blog = await download_small_avatar(user.id, context) + avatar_big_blog = await common.download_big_avatar(user.id, context) + avatar_small_blog = await common.download_small_avatar(user.id, context) avatar_big_id = None if avatar_big_blog: sent_message = await update.effective_chat.send_photo( @@ -95,14 +81,14 @@ async def user_data_refresh(update: Update, context: ContextTypes.DEFAULT_TYPE): ) avatar_big_id = sent_message.photo[-1].file_id await sent_message.delete() - db_user = get_user_by_id(user.id) + db_user = dao.get_user_by_id(user.id) db_user.username = username db_user.full_name = full_name db_user.avatar_big_blob = avatar_big_blog db_user.avatar_small_blob = avatar_small_blog db_user.avatar_big_id = avatar_big_id - db.commit() - info = get_user_info(user) + dao.commit() + info = common.get_user_info(user) info += "\n刷新成功" if avatar_big_id: await query.edit_message_media( @@ -125,7 +111,7 @@ async def user_waifu_manage(update: Update, context: ContextTypes.DEFAULT_TYPE): if "divorce" in query_data: await _divorce_ask(update, context) return - db_user = get_user_by_id(update.effective_user.id) + db_user = dao.get_user_by_id(update.effective_user.id) if "set_waifu_mention" in query_data: db_user.waifu_mention = not db_user.waifu_mention set_mention_text = "别@你" if db_user.waifu_mention else "抽到你时@你" @@ -140,9 +126,9 @@ async def user_waifu_manage(update: Update, context: ContextTypes.DEFAULT_TYPE): [InlineKeyboardButton(text="返回", callback_data="back_home")], ] ) - text = get_user_waifu_info(update.effective_user) + text = common.get_user_waifu_info(update.effective_user) await query.edit_message_caption(caption=text, reply_markup=waifu_manage_markup) - db.commit() + dao.commit() _divorce_ask_markup = InlineKeyboardMarkup( @@ -160,11 +146,11 @@ async def _divorce_ask(update: Update, context: ContextTypes.DEFAULT_TYPE): if "divorce_confirm" in query.data: await _divorce_confirm(update, context) return - db_user = get_user_by_id(update.effective_user.id) + db_user = dao.get_user_by_id(update.effective_user.id) if not db_user.is_married: await query.answer("可是...你还没有结婚呀qwq", show_alert=True, cache_time=15) return - married_waifu = get_user_by_id(db_user.married_waifu_id) + married_waifu = dao.get_user_by_id(db_user.married_waifu_id) if query.message.photo and married_waifu.avatar_big_id: await query.edit_message_media( @@ -182,20 +168,20 @@ async def _divorce_ask(update: Update, context: ContextTypes.DEFAULT_TYPE): async def _divorce_confirm(update: Update, context: ContextTypes.DEFAULT_TYPE): - db_user = get_user_by_id(update.effective_user.id) + db_user = dao.get_user_by_id(update.effective_user.id) query = update.callback_query - married_waifu = get_user_by_id(db_user.married_waifu_id) + married_waifu = dao.get_user_by_id(db_user.married_waifu_id) db_user.is_married = False db_user.married_waifu_id = None married_waifu.is_married = False married_waifu.married_waifu_id = None - db.commit() - refresh_user_all_waifu(db_user) - refresh_user_all_waifu(married_waifu) + dao.commit() + dao.refresh_user_all_waifu(db_user) + dao.refresh_user_all_waifu(married_waifu) await query.edit_message_caption( caption="_愿你有一天和重要之人重逢_", parse_mode="MarkdownV2", - reply_markup=back_home_markup, + reply_markup=common.back_home_markup, ) logger.debug( f"{db_user.full_name}<{db_user.id}> divorced {married_waifu.full_name}<{married_waifu.id}>" # noqa: E501 @@ -208,7 +194,7 @@ async def delete_user_quote(update: Update, context: ContextTypes.DEFAULT_TYPE): await _user_quote_manage(update, context) return quote_link = query.data.split(" ")[1] - delete_quote_by_link(quote_link) + dao.delete_quote_by_link(quote_link) await query.answer("已删除", show_alert=False, cache_time=5) await _user_quote_manage(update, context) @@ -219,7 +205,7 @@ async def _user_quote_manage(update: Update, context: ContextTypes.DEFAULT_TYPE) logger.info(f"({user.name}) ") page = int(query.data.split(" ")[-1]) if len(query.data.split(" ")) > 1 else 1 page_size = 5 - quotes_count = get_user_quotes_count(user) + quotes_count = dao.get_user_quotes_count(user) max_page = ceil(quotes_count / page_size) if quotes_count == 0: caption = ( @@ -227,7 +213,7 @@ async def _user_quote_manage(update: Update, context: ContextTypes.DEFAULT_TYPE) ) # noqa: E501 await query.edit_message_caption( caption=caption, - reply_markup=back_home_markup, + reply_markup=common.back_home_markup, ) return if page > max_page or page < 1: @@ -235,7 +221,7 @@ async def _user_quote_manage(update: Update, context: ContextTypes.DEFAULT_TYPE) return text = f"你的语录: 共{quotes_count}条; 第{page}/{max_page}页\n" text += "点击序号删除语录\n\n" - quotes = get_user_quotes_page(user, page, page_size) + quotes = dao.get_user_quotes_page(user, page, page_size) keyboard, line = [], [] for index, quote in enumerate(quotes): quote_content = ( diff --git a/src/callbacks/waifu.py b/kmua/callbacks/waifu.py similarity index 78% rename from src/callbacks/waifu.py rename to kmua/callbacks/waifu.py index 97bf505..ae7468a 100644 --- a/src/callbacks/waifu.py +++ b/kmua/callbacks/waifu.py @@ -6,47 +6,10 @@ from telegram.ext import ContextTypes from telegram.helpers import escape_markdown -from ..common.message import message_recorder -from ..common.user import ( - download_small_avatar, - fake_users_id, - get_big_avatar_bytes, - mention_markdown_v2, - verify_user_can_manage_bot_in_chat, -) -from ..common.utils import loading_word -from ..common.waifu import ( - get_chat_waifu_relationships, - get_marry_markup, - get_remove_markup, - get_waifu_markup, - get_waifu_text, - render_waifu_graph, -) -from ..dao import db -from ..dao.association import ( - delete_association_in_chat, -) -from ..dao.chat import ( - get_chat_users_without_bots_id, -) -from ..dao.user import ( - add_user, - get_user_by_id, -) -from ..logger import logger -from ..models.models import ChatData, UserData -from ..service.user import ( - check_user_in_chat, -) -from ..service.waifu import ( - get_chat_married_users_id, - get_chat_user_participated_waifu, - get_user_waifu_in_chat, - put_user_waifu_in_chat, - refresh_user_all_waifu, - refresh_user_waifu_in_chat, -) +from kmua.logger import logger +import kmua.common as common +import kmua.dao as dao +from kmua.models import ChatData, UserData async def waifu_graph(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -74,8 +37,8 @@ async def send_waifu_graph( ): logger.debug(f"Generating waifu graph for {chat.title}<{chat.id}>") try: - relationships = get_chat_waifu_relationships(chat) - participate_users = get_chat_user_participated_waifu(chat) + relationships = common.get_chat_waifu_relationships(chat) + participate_users = dao.get_chat_user_participated_waifu(chat) if not participate_users or not relationships: await context.bot.send_message( chat.id, @@ -86,7 +49,7 @@ async def send_waifu_graph( await context.bot.send_message( chat.id, - random.choice(loading_word), + random.choice(common.loading_word), reply_to_message_id=msg_id, ) @@ -97,7 +60,7 @@ async def send_waifu_graph( } for user in participate_users } - image_bytes = render_waifu_graph(relationships, user_info) + image_bytes = common.render_waifu_graph(relationships, user_info) logger.debug(f"image_size: {len(image_bytes)}") await context.bot.send_document( chat.id, @@ -136,7 +99,7 @@ async def today_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE): context.user_data["waifu_waiting"] = True waifu: UserData = None - message_recorder(update, context) + common.message_recorder(update, context) try: await context.bot.send_chat_action(chat.id, ChatAction.TYPING) waifu, is_got_waifu = await _get_waifu_for_user(update, context, user, chat) @@ -145,16 +108,16 @@ async def today_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE): if waifu.is_married and user.id != waifu.married_waifu_id: await update.effective_message.reply_text("你没能抽到老婆, 再试一次吧~") return - is_waifu_in_chat = check_user_in_chat(waifu, chat) + is_waifu_in_chat = dao.check_user_in_chat(waifu, chat) if is_waifu_in_chat: - put_user_waifu_in_chat(user, chat, waifu) - waifu_markup = get_waifu_markup(waifu, user) - text = get_waifu_text(waifu, is_got_waifu) + dao.put_user_waifu_in_chat(user, chat, waifu) + waifu_markup = common.get_waifu_markup(waifu, user) + text = common.get_waifu_text(waifu, is_got_waifu) if user.id == waifu.married_waifu_id: text = f"你和 [{escape_markdown(waifu.full_name,2)}](tg://user?id={waifu.id}) 已经结婚了哦, 还想娶第二遍嘛?" # noqa: E501 waifu_markup = None if is_waifu_in_chat: - put_user_waifu_in_chat(waifu, chat, user) + dao.put_user_waifu_in_chat(waifu, chat, user) photo_to_send = await _get_photo_to_send(waifu, context) if photo_to_send is None: @@ -185,11 +148,15 @@ async def today_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE): finally: if waifu: if not waifu.avatar_small_blob: - waifu.avatar_small_blob = await download_small_avatar(waifu.id, context) - db_user = get_user_by_id(user.id) + waifu.avatar_small_blob = await common.download_small_avatar( + waifu.id, context + ) + db_user = dao.get_user_by_id(user.id) if not db_user.avatar_small_blob: - db_user.avatar_small_blob = await download_small_avatar(user.id, context) - db.commit() + db_user.avatar_small_blob = await common.download_small_avatar( + user.id, context + ) + dao.commit() context.user_data["waifu_waiting"] = False @@ -201,7 +168,7 @@ async def _get_waifu_for_user( :return: (waifu, is_got_waifu) """ - if waifu := get_user_waifu_in_chat(user, chat): + if waifu := dao.get_user_waifu_in_chat(user, chat): return waifu, True group_member = await _get_chat_members_id_to_get_waifu(update, context, user, chat) if not group_member: @@ -209,11 +176,11 @@ async def _get_waifu_for_user( waifu_id = random.choice(group_member) while retry := 0 < 3: try: - if waifu := get_user_by_id(waifu_id): + if waifu := dao.get_user_by_id(waifu_id): return waifu, False else: waifu_chat = await context.bot.get_chat(waifu_id) - waifu = add_user(waifu_chat) + waifu = dao.add_user(waifu_chat) return waifu, False except Exception as e: logger.error( @@ -229,9 +196,9 @@ async def _get_waifu_for_user( async def _get_chat_members_id_to_get_waifu( update: Update, context: ContextTypes.DEFAULT_TYPE, user: User, chat: Chat ) -> list[int]: - group_member = get_chat_users_without_bots_id(chat) - married = get_chat_married_users_id(chat) - to_remove = set(married + fake_users_id + [user.id]) + group_member = dao.get_chat_users_without_bots_id(chat) + married = dao.get_chat_married_users_id(chat) + to_remove = set(married + common.fake_users_id + [user.id]) group_member = [i for i in group_member if i not in to_remove] if not group_member: await update.message.reply_text(text="你现在没有老婆, 因为咱的记录中找不到其他群友") # noqa: E501 @@ -244,7 +211,7 @@ async def _get_photo_to_send( ) -> bytes | str | None: if waifu.avatar_big_id: return waifu.avatar_big_id - return await get_big_avatar_bytes(waifu.id, context) + return await common.get_big_avatar_bytes(waifu.id, context) async def remove_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -256,19 +223,19 @@ async def remove_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE): await _remove_waifu_cancel(update, context) return query_data = query_data.split(" ") - if not await verify_user_can_manage_bot_in_chat( + if not await common.verify_user_can_manage_bot_in_chat( update.effective_user, update.effective_chat, update, context ): return waifu_id = int(query_data[1]) user_id = int(query_data[2]) message = update.callback_query.message - user = get_user_by_id(user_id) - waifu = get_user_by_id(waifu_id) + user = dao.get_user_by_id(user_id) + waifu = dao.get_user_by_id(waifu_id) if not user or not waifu: await update.callback_query.answer(text="查无此人,可能已经被移除了") return - markup = get_remove_markup(waifu, user) + markup = common.get_remove_markup(waifu, user) if message.photo: await message.edit_caption(caption="确定要移除ta吗?", reply_markup=markup) return @@ -276,7 +243,7 @@ async def remove_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE): async def _remove_waifu_confirm(update: Update, context: ContextTypes.DEFAULT_TYPE): - if not await verify_user_can_manage_bot_in_chat( + if not await common.verify_user_can_manage_bot_in_chat( update.effective_user, update.effective_chat, update, context ): return @@ -284,13 +251,13 @@ async def _remove_waifu_confirm(update: Update, context: ContextTypes.DEFAULT_TY waifu_id = int(query_data[1]) user_id = int(query_data[2]) message = update.callback_query.message - user = get_user_by_id(user_id) - waifu = get_user_by_id(waifu_id) + user = dao.get_user_by_id(user_id) + waifu = dao.get_user_by_id(waifu_id) if not user or not waifu: await update.callback_query.answer(text="查无此人,可能已经被移除了") return - delete_association_in_chat(message.chat, waifu) - refresh_user_waifu_in_chat(user, message.chat) + dao.delete_association_in_chat(message.chat, waifu) + dao.refresh_user_waifu_in_chat(user, message.chat) text = ( f"_已移除该用户:_ {escape_markdown(waifu.full_name,2)}\n" + f"ta 曾是 {escape_markdown(user.full_name,2)} 的老婆" @@ -305,7 +272,7 @@ async def _remove_waifu_confirm(update: Update, context: ContextTypes.DEFAULT_TY async def _remove_waifu_cancel(update: Update, context: ContextTypes.DEFAULT_TYPE): - if not await verify_user_can_manage_bot_in_chat( + if not await common.verify_user_can_manage_bot_in_chat( update.effective_user, update.effective_chat, update, context ): return @@ -313,10 +280,10 @@ async def _remove_waifu_cancel(update: Update, context: ContextTypes.DEFAULT_TYP waifu_id = int(query_data[1]) user_id = int(query_data[2]) message = update.callback_query.message - user = get_user_by_id(user_id) - waifu = get_user_by_id(waifu_id) - text = get_waifu_text(waifu, False) - markup = get_waifu_markup(waifu, user) + user = dao.get_user_by_id(user_id) + waifu = dao.get_user_by_id(waifu_id) + text = common.get_waifu_text(waifu, False) + markup = common.get_waifu_markup(waifu, user) if message.photo: await message.edit_caption( caption=text, @@ -346,8 +313,8 @@ async def marry_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE): query_data = query.data.split(" ") waifu_id = int(query_data[1]) user_id = int(query_data[2]) - db_waifu = get_user_by_id(waifu_id) - db_user = get_user_by_id(user_id) + db_waifu = dao.get_user_by_id(waifu_id) + db_user = dao.get_user_by_id(user_id) message = query.message if not db_waifu.is_real_user: await query.answer( @@ -375,10 +342,10 @@ async def marry_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE): cache_time=60, ) return - waifu_mention = mention_markdown_v2(db_waifu) - user_mention = mention_markdown_v2(db_user) + waifu_mention = common.mention_markdown_v2(db_waifu) + user_mention = common.mention_markdown_v2(db_user) text = f"{waifu_mention}, 你愿意和 {user_mention} 结婚吗qwq?" - reply_markup = get_marry_markup(waifu_id, user_id) + reply_markup = common.get_marry_markup(waifu_id, user_id) if message.photo: await message.edit_caption( caption=text, @@ -399,8 +366,8 @@ async def _agree_marry_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE) query_data = query.data.split(" ") waifu_id = int(query_data[1]) user_id = int(query_data[2]) - db_waifu = get_user_by_id(waifu_id) - db_user = get_user_by_id(user_id) + db_waifu = dao.get_user_by_id(waifu_id) + db_user = dao.get_user_by_id(user_id) message = query.message if now_user.id != waifu_id: await query.answer( @@ -418,15 +385,15 @@ async def _agree_marry_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE) db_user.is_married = True db_waifu.married_waifu_id = user_id db_waifu.is_married = True - db.commit() - db_user = get_user_by_id(user_id) - db_waifu = get_user_by_id(waifu_id) - refresh_user_all_waifu(db_user) - refresh_user_all_waifu(db_waifu) - put_user_waifu_in_chat(db_user, message.chat, db_waifu) - put_user_waifu_in_chat(db_waifu, message.chat, db_user) - waifu_mention = mention_markdown_v2(db_waifu) - user_mention = mention_markdown_v2(db_user) + dao.commit() + db_user = dao.get_user_by_id(user_id) + db_waifu = dao.get_user_by_id(waifu_id) + dao.refresh_user_all_waifu(db_user) + dao.refresh_user_all_waifu(db_waifu) + dao.put_user_waifu_in_chat(db_user, message.chat, db_waifu) + dao.put_user_waifu_in_chat(db_waifu, message.chat, db_user) + waifu_mention = common.mention_markdown_v2(db_waifu) + user_mention = common.mention_markdown_v2(db_user) text = f"恭喜 {waifu_mention} 和 {user_mention} 结婚啦\~" if message.photo: await message.edit_caption( @@ -445,14 +412,14 @@ async def _refuse_marry_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE query = update.callback_query query_data = query.data.split(" ") waifu_id = int(query_data[1]) - db_waifu = get_user_by_id(waifu_id) + db_waifu = dao.get_user_by_id(waifu_id) message = query.message if now_user.id != waifu_id: await query.answer( text="(￣ε(#￣) 别人的事情咱不要打扰呢", show_alert=True, cache_time=60 ) # noqa: E501 return - text = get_waifu_text(db_waifu, False) + text = common.get_waifu_text(db_waifu, False) await query.answer( text="(´。＿。｀) 你拒绝了 ta 的求婚呢. 有些人一旦错过就不再...", show_alert=True, @@ -482,14 +449,14 @@ async def _cancel_marry_waifu(update: Update, context: ContextTypes.DEFAULT_TYPE text="(￣ε(#￣) 别人的事情咱不要打扰呢", show_alert=True, cache_time=60 ) # noqa: E501 return - db_waifu = get_user_by_id(waifu_id) - db_user = get_user_by_id(user_id) - text = get_waifu_text(db_waifu, False) + db_waifu = dao.get_user_by_id(waifu_id) + db_user = dao.get_user_by_id(user_id) + text = common.get_waifu_text(db_waifu, False) await query.answer( text="o((>ω< ))o 你取消了这个求婚请求", cache_time=5, ) - markup = get_waifu_markup(db_waifu, db_user) + markup = common.get_waifu_markup(db_waifu, db_user) if message.photo: await message.edit_caption( caption=text, diff --git a/kmua/common/__init__.py b/kmua/common/__init__.py new file mode 100644 index 0000000..12a69fe --- /dev/null +++ b/kmua/common/__init__.py @@ -0,0 +1,7 @@ +from .bot import * # noqa +from .chat import * # noqa +from .user import * # noqa +from .utils import * # noqa +from .message import * # noqa +from .quote import * # noqa +from .waifu import * # noqa diff --git a/src/common/bot.py b/kmua/common/bot.py similarity index 70% rename from src/common/bot.py rename to kmua/common/bot.py index 0e3bb08..c2355b3 100644 --- a/src/common/bot.py +++ b/kmua/common/bot.py @@ -5,11 +5,8 @@ import psutil -from ..config import data_dir, settings -from ..dao.association import get_all_associations_count -from ..dao.chat import get_all_chats_count -from ..dao.quote import get_all_quotes_count -from ..dao.user import get_all_users_count +import kmua.dao as dao +from kmua.config import data_dir, settings _database_path = Path( data_dir / settings.get("db_url", "sqlite:///./data/kmua.db").split("/")[-1] @@ -19,10 +16,10 @@ def get_bot_status() -> str: db_status = f""" Database Status: - - Users: {get_all_users_count()} - - Chats: {get_all_chats_count()} - - Quotes: {get_all_quotes_count()} - - Associations: {get_all_associations_count()} + - Users: {dao.get_all_users_count()} + - Chats: {dao.get_all_chats_count()} + - Quotes: {dao.get_all_quotes_count()} + - Associations: {dao.get_all_associations_count()} - Size: {_database_path.stat().st_size / 1024 / 1024:.2f} MB """ pid = os.getpid() diff --git a/src/common/chat.py b/kmua/common/chat.py similarity index 82% rename from src/common/chat.py rename to kmua/common/chat.py index dd87d33..23cd588 100644 --- a/src/common/chat.py +++ b/kmua/common/chat.py @@ -1,11 +1,11 @@ from telegram import Chat -from ..dao.chat import add_chat -from ..models.models import ChatData +import kmua.dao as dao +from kmua.models import ChatData def get_chat_info(chat: Chat | ChatData) -> str: - db_chat = add_chat(chat) + db_chat = dao.add_chat(chat) text = f"chat_id: {db_chat.id}\n" text += f"title: {db_chat.title}\n\n" text += f"记录中共有 {len(db_chat.members)} 个成员\n" diff --git a/src/common/message.py b/kmua/common/message.py similarity index 87% rename from src/common/message.py rename to kmua/common/message.py index d8c596e..b4a3314 100644 --- a/src/common/message.py +++ b/kmua/common/message.py @@ -2,9 +2,8 @@ from telegram.constants import ChatID, ChatType from telegram.ext import ContextTypes -from ..dao.association import add_association_in_chat -from ..dao.user import add_user -from ..logger import logger +import kmua.dao as dao +from kmua.logger import logger def message_recorder(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -21,9 +20,9 @@ def message_recorder(update: Update, context: ContextTypes.DEFAULT_TYPE): return if message.sender_chat: user = message.sender_chat - db_user = add_user(user) + db_user = dao.add_user(user) if chat.type == ChatType.GROUP or chat.type == ChatType.SUPERGROUP: - add_association_in_chat(chat, db_user) + dao.add_association_in_chat(chat, db_user) def get_message_common_link(message: Message) -> str: diff --git a/src/common/quote.py b/kmua/common/quote.py similarity index 88% rename from src/common/quote.py rename to kmua/common/quote.py index 4459589..7f6badb 100644 --- a/src/common/quote.py +++ b/kmua/common/quote.py @@ -8,15 +8,11 @@ async def generate_quote_img(avatar: bytes, text: str, name: str) -> bytes: text = text.replace("\n", " ") - font_path = str( - Path(__file__).resolve().parent.parent.parent / "resource" / "TsukuA.ttc" - ) + font_path = str(Path(__file__).resolve().parent.parent / "resource" / "TsukuA.ttc") font_size = 42 font = ImageFont.truetype(font_path, font_size) base_img = Image.open( - os.path.join( - Path(os.path.dirname(__file__)).parent.parent, "resource", "base.png" - ) + os.path.join(Path(os.path.dirname(__file__)).parent, "resource", "base.png") ) img = Image.new("RGBA", (1200, 630), (255, 255, 255, 0)) avatar = Image.open(io.BytesIO(avatar)) diff --git a/src/common/user.py b/kmua/common/user.py similarity index 92% rename from src/common/user.py rename to kmua/common/user.py index 6285e7a..9bc8c1c 100644 --- a/src/common/user.py +++ b/kmua/common/user.py @@ -7,12 +7,11 @@ from telegram.ext import ContextTypes from telegram.helpers import escape_markdown -from ..config import settings -from ..dao import db -from ..dao.user import add_user, get_user_by_id -from ..logger import logger -from ..models.models import ChatData, UserData -from ..service.user import get_user_is_bot_admin_in_chat +from kmua.config import settings +import kmua.dao as dao +from kmua.logger import logger +from kmua.models import ChatData, UserData + fake_users_id = [ChatID.FAKE_CHANNEL, ChatID.ANONYMOUS_ADMIN, ChatID.SERVICE_CHAT] @@ -21,7 +20,7 @@ async def get_big_avatar_bytes( chat_id: int, context: ContextTypes.DEFAULT_TYPE ) -> bytes | None: logger.debug(f"Get big avatar for {chat_id}") - db_user = get_user_by_id(chat_id) + db_user = dao.get_user_by_id(chat_id) if db_user: if db_user.avatar_big_blob: return db_user.avatar_big_blob @@ -29,7 +28,7 @@ async def get_big_avatar_bytes( avatar = await download_big_avatar(chat_id, context) if avatar: db_user.avatar_big_blob = avatar - db.commit() + dao.commit() return avatar else: return await download_big_avatar(chat_id, context) @@ -56,7 +55,7 @@ async def get_small_avatar_bytes( chat_id: int, context: ContextTypes.DEFAULT_TYPE ) -> bytes | None: logger.debug(f"Get small avatar for {chat_id}") - db_user = get_user_by_id(chat_id) + db_user = dao.get_user_by_id(chat_id) if db_user: if db_user.avatar_small_blob: return db_user.avatar_small_blob @@ -64,7 +63,7 @@ async def get_small_avatar_bytes( avatar = await download_small_avatar(chat_id, context) if avatar: db_user.avatar_small_blob = avatar - db.commit() + dao.commit() return avatar else: return await download_small_avatar(chat_id, context) @@ -144,7 +143,7 @@ def verify_user_can_manage_bot(user: User | UserData) -> bool: """ if user.id in settings.owners: return True - if db_user := get_user_by_id(user.id): + if db_user := dao.get_user_by_id(user.id): return db_user.is_bot_global_admin return False @@ -166,7 +165,7 @@ async def verify_user_can_manage_bot_in_chat( if ( verify_user_can_manage_bot(user) or user.id == ChatID.ANONYMOUS_ADMIN - or get_user_is_bot_admin_in_chat(user, chat) + or dao.get_user_is_bot_admin_in_chat(user, chat) ): return True if await verify_user_is_chat_owner(user, chat, update, context): @@ -176,7 +175,7 @@ async def verify_user_can_manage_bot_in_chat( def get_user_info(user: User | UserData) -> str: logger.debug(f"Get user info for {user.full_name}<{user.id}>") - db_user = add_user(user) + db_user = dao.add_user(user) info = f""" id: {db_user.id} username: {db_user.username} @@ -194,5 +193,5 @@ def get_user_info(user: User | UserData) -> str: def mention_markdown_v2(user: User | UserData | Chat | ChatData) -> str: - db_user = add_user(user) + db_user = dao.add_user(user) return f"[{escape_markdown(db_user.full_name,2)}](tg://user?id={db_user.id})" diff --git a/src/common/utils.py b/kmua/common/utils.py similarity index 97% rename from src/common/utils.py rename to kmua/common/utils.py index 67249ea..943c3e7 100644 --- a/src/common/utils.py +++ b/kmua/common/utils.py @@ -4,7 +4,7 @@ import re from operator import attrgetter import glob -from ..logger import logger +from kmua.logger import logger import json from telegram import ( @@ -55,7 +55,7 @@ def parse_arguments(text: str) -> list[str]: def _load_word_dict(): word_dict_path_internal = ( - pathlib.Path(__file__).parent.parent.parent / "resource" / "word_dicts" + pathlib.Path(__file__).parent.parent / "resource" / "word_dicts" ) word_dict_path_user = ( pathlib.Path(__file__).parent.parent.parent / "data" / "word_dicts" diff --git a/src/common/waifu.py b/kmua/common/waifu.py similarity index 92% rename from src/common/waifu.py rename to kmua/common/waifu.py index d102c31..60a7d09 100644 --- a/src/common/waifu.py +++ b/kmua/common/waifu.py @@ -7,15 +7,9 @@ from telegram import Chat, InlineKeyboardButton, InlineKeyboardMarkup, User from telegram.helpers import escape_markdown -from ..dao.chat import get_chat_members -from ..dao.user import add_user -from ..logger import logger -from ..models.models import ChatData, UserData -from ..service.waifu import ( - get_user_waifu_in_chat_exclude_married, - get_user_waifus_of_with_chat, - get_user_waifus_with_chat, -) +import kmua.dao as dao +from kmua.logger import logger +from kmua.models import ChatData, UserData def get_chat_waifu_relationships( @@ -23,9 +17,9 @@ def get_chat_waifu_relationships( ) -> list[tuple[int, int]]: relationships = [] logger.debug(f"Get chat waifu relationships for {chat.title}<{chat.id}>") - members = get_chat_members(chat) + members = dao.get_chat_members(chat) for member in members: - waifu = get_user_waifu_in_chat_exclude_married(member, chat) + waifu = dao.get_user_waifu_in_chat_exclude_married(member, chat) if waifu: relationships.append((member.id, waifu.id)) return relationships @@ -44,13 +38,13 @@ def get_chat_waifu_info_dict( def get_user_waifu_info(user: User | UserData) -> str: logger.debug(f"Get user waifu info for {user.full_name}<{user.id}>") - db_user = add_user(user) + db_user = dao.add_user(user) text = f""" 是否@你: {db_user.waifu_mention} 已婚: {db_user.is_married} 已婚老婆id: {db_user.married_waifu_id} """ - waifus_with_chat = get_user_waifus_with_chat(user) + waifus_with_chat = dao.get_user_waifus_with_chat(user) if not waifus_with_chat: text += "\n你今天还没有老婆哦" else: @@ -59,7 +53,7 @@ def get_user_waifu_info(user: User | UserData) -> str: text += "\n你今天的老婆们(最多显示233条):\n" for waifu, chat in waifus_with_chat: text += f"{waifu.full_name} ({chat.title})\n" - waifus_of_with_chat = get_user_waifus_of_with_chat(user) + waifus_of_with_chat = dao.get_user_waifus_of_with_chat(user) if not waifus_of_with_chat: text += "\n今天还没有人把你当老婆哦" else: diff --git a/src/config/__init__.py b/kmua/config.py similarity index 74% rename from src/config/__init__.py rename to kmua/config.py index 8371458..789cdb3 100644 --- a/src/config/__init__.py +++ b/kmua/config.py @@ -2,9 +2,9 @@ from dynaconf import Dynaconf -_BASE_DIR = Path(__file__).resolve().parent.parent.parent +_BASE_DIR = Path(__file__).resolve().parent.parent -data_dir = Path(f"{Path(__file__).resolve().parent.parent.parent}/data") +data_dir = Path(f"{Path(__file__).resolve().parent.parent}/data") settings_files = [ "settings.toml", diff --git a/kmua/dao/__init__.py b/kmua/dao/__init__.py new file mode 100644 index 0000000..08b2fbc --- /dev/null +++ b/kmua/dao/__init__.py @@ -0,0 +1,7 @@ +from .association import * # noqa: F403 +from .chat import * # noqa: F403 +from .quote import * # noqa: F403 +from .user import * # noqa: F403 +from .chat_service import * # noqa: F403 +from .user_service import * # noqa: F403 +from .waifu import * # noqa: F403 diff --git a/src/dao/db.py b/kmua/dao/_db.py similarity index 53% rename from src/dao/db.py rename to kmua/dao/_db.py index 894a3d9..8d96f3f 100644 --- a/src/dao/db.py +++ b/kmua/dao/_db.py @@ -1,21 +1,20 @@ from sqlalchemy import create_engine -from sqlalchemy.orm import declarative_base, scoped_session, sessionmaker +from sqlalchemy.orm import scoped_session, sessionmaker -from ..config import settings +from kmua.config import settings engine = create_engine(settings.DB_URL) _session = scoped_session(sessionmaker(autoflush=False, bind=engine)) -db = _session() -Base = declarative_base() +_db = _session() def commit(): try: - db.commit() + _db.commit() except Exception as err: - db.rollback() + _db.rollback() raise err def close(): - db.close() + _db.close() diff --git a/src/dao/association.py b/kmua/dao/association.py similarity index 76% rename from src/dao/association.py rename to kmua/dao/association.py index 2e658a4..ed9148e 100644 --- a/src/dao/association.py +++ b/kmua/dao/association.py @@ -1,12 +1,12 @@ from telegram import Chat, User -from ..models.models import ChatData, UserChatAssociation, UserData -from .db import db, commit +from kmua.dao._db import commit, _db +from kmua.models import ChatData, UserChatAssociation, UserData def get_association_in_chat(chat: Chat | ChatData) -> UserChatAssociation | None: return ( - db.query(UserChatAssociation) + _db.query(UserChatAssociation) .filter(UserChatAssociation.chat_id == chat.id) .all() ) @@ -16,7 +16,7 @@ def get_association_in_chat_by_user( chat: Chat | ChatData, user: User | UserData | Chat | ChatData ) -> UserChatAssociation | None: return ( - db.query(UserChatAssociation) + _db.query(UserChatAssociation) .filter(UserChatAssociation.chat_id == chat.id) .filter(UserChatAssociation.user_id == user.id) .first() @@ -37,7 +37,7 @@ def add_association_in_chat( """ if association := get_association_in_chat_by_user(chat, user): return association - db.add( + _db.add( UserChatAssociation( user_id=user.id, chat_id=chat.id, @@ -58,32 +58,34 @@ def delete_association_in_chat( :param chat: Chat or ChatData object """ if association := get_association_in_chat_by_user(chat, user): - db.delete(association) + _db.delete(association) commit() def get_associations_of_user(user: User | UserData) -> list[UserChatAssociation]: - return db.query(UserChatAssociation).filter_by(user_id=user.id).all() + return _db.query(UserChatAssociation).filter_by(user_id=user.id).all() def get_associations_of_user_waifu_of_in_chat( user: User | UserData, chat: Chat | ChatData ) -> list[UserChatAssociation]: return ( - db.query(UserChatAssociation).filter_by(waifu_id=user.id, chat_id=chat.id).all() + _db.query(UserChatAssociation) + .filter_by(waifu_id=user.id, chat_id=chat.id) + .all() ) def get_associations_of_user_waifu_of( user: User | UserData, ) -> list[UserChatAssociation]: - return db.query(UserChatAssociation).filter_by(waifu_id=user.id).all() + return _db.query(UserChatAssociation).filter_by(waifu_id=user.id).all() def update_associations_all_waifu_id_to_none(): - db.query(UserChatAssociation).update({UserChatAssociation.waifu_id: None}) - db.commit() + _db.query(UserChatAssociation).update({UserChatAssociation.waifu_id: None}) + _db.commit() def get_all_associations_count() -> int: - return db.query(UserChatAssociation).count() + return _db.query(UserChatAssociation).count() diff --git a/src/dao/chat.py b/kmua/dao/chat.py similarity index 69% rename from src/dao/chat.py rename to kmua/dao/chat.py index 5c268f1..2e19a85 100644 --- a/src/dao/chat.py +++ b/kmua/dao/chat.py @@ -1,11 +1,11 @@ from telegram import Chat -from ..models.models import ChatData, Quote, UserData -from .db import db, commit +from kmua.dao._db import commit, _db +from kmua.models import ChatData, Quote, UserData def get_chat_by_id(chat_id: int) -> ChatData | None: - return db.query(ChatData).filter(ChatData.id == chat_id).first() + return _db.query(ChatData).filter(ChatData.id == chat_id).first() def add_chat(chat: Chat | ChatData) -> ChatData: @@ -17,7 +17,7 @@ def add_chat(chat: Chat | ChatData) -> ChatData: """ if chatdata := get_chat_by_id(chat.id): return chatdata - db.add( + _db.add( ChatData( id=chat.id, title=chat.title, @@ -28,11 +28,11 @@ def add_chat(chat: Chat | ChatData) -> ChatData: def get_chat_members(chat: Chat | ChatData) -> list[UserData]: - db_chat = get_chat_by_id(chat.id) - if db_chat is None: + _db_chat = get_chat_by_id(chat.id) + if _db_chat is None: add_chat(chat) return [] - return db_chat.members + return _db_chat.members def get_chat_members_id(chat: Chat | ChatData) -> list[int]: @@ -41,28 +41,28 @@ def get_chat_members_id(chat: Chat | ChatData) -> list[int]: def get_chat_quote_probability(chat: Chat | ChatData) -> float: - db_chat = get_chat_by_id(chat.id) - if db_chat is None: + _db_chat = get_chat_by_id(chat.id) + if _db_chat is None: add_chat(chat) return 0.001 - return db_chat.quote_probability + return _db_chat.quote_probability def update_chat_quote_probability(chat: Chat | ChatData, probability: float): - db_chat = get_chat_by_id(chat.id) - if db_chat is None: + _db_chat = get_chat_by_id(chat.id) + if _db_chat is None: add_chat(chat) - db_chat = get_chat_by_id(chat.id) - db_chat.quote_probability = probability + _db_chat = get_chat_by_id(chat.id) + _db_chat.quote_probability = probability commit() def get_chat_quotes(chat: Chat | ChatData) -> list[Quote]: - db_chat = get_chat_by_id(chat.id) - if db_chat is None: + _db_chat = get_chat_by_id(chat.id) + if _db_chat is None: add_chat(chat) return [] - return db_chat.quotes + return _db_chat.quotes def get_chat_quotes_count(chat: Chat | ChatData) -> int: @@ -84,8 +84,8 @@ def get_chat_bots(chat: Chat | ChatData) -> list[UserData]: """ 获取 chat 中的 bot """ - db_chat = add_chat(chat) - return [user for user in db_chat.members if user.is_bot] + _db_chat = add_chat(chat) + return [user for user in _db_chat.members if user.is_bot] def get_chat_bots_id(chat: Chat | ChatData) -> list[int]: @@ -94,8 +94,8 @@ def get_chat_bots_id(chat: Chat | ChatData) -> list[int]: def get_chat_users_without_bots(chat: Chat | ChatData) -> list[UserData]: - db_chat = add_chat(chat) - return [user for user in db_chat.members if not user.is_bot] + _db_chat = add_chat(chat) + return [user for user in _db_chat.members if not user.is_bot] def get_chat_users_without_bots_id(chat: Chat | ChatData) -> list[int]: @@ -104,7 +104,7 @@ def get_chat_users_without_bots_id(chat: Chat | ChatData) -> list[int]: def get_all_chats() -> list[ChatData]: - return db.query(ChatData).all() + return _db.query(ChatData).all() def get_all_chats_id() -> list[int]: @@ -112,12 +112,12 @@ def get_all_chats_id() -> list[int]: def delete_chat(chat: Chat | ChatData): - db_chat = get_chat_by_id(chat.id) - if db_chat is None: + _db_chat = get_chat_by_id(chat.id) + if _db_chat is None: return - db.delete(db_chat) + _db.delete(_db_chat) commit() def get_all_chats_count() -> int: - return db.query(ChatData).count() + return _db.query(ChatData).count() diff --git a/kmua/dao/chat_service.py b/kmua/dao/chat_service.py new file mode 100644 index 0000000..d20ba16 --- /dev/null +++ b/kmua/dao/chat_service.py @@ -0,0 +1,40 @@ +from telegram import Chat +import kmua.dao.chat as chat_dao +import kmua.dao.quote as quote_dao +from ._db import commit +from kmua.models import ChatData + + +def delete_chat_data_and_quotes(chat: Chat | ChatData): + # 删除与该群组相关的所有数据 + db_chat = chat_dao.get_chat_by_id(chat.id) + if db_chat is None: + return + for quote in db_chat.quotes: + quote_dao.delete_quote(quote) + chat_dao.delete_chat(db_chat) + commit() + + +def update_chat_id(old_id: int, new_id: int): + db_chat = chat_dao.get_chat_by_id(old_id) + if db_chat is None: + return + if chat_dao.get_chat_by_id(new_id) is not None: + return + db_chat.id = new_id + commit() + + +def update_chat_title(chat: Chat | ChatData, title: str): + db_chat = chat_dao.get_chat_by_id(chat.id) + if db_chat is None: + return + db_chat.title = title + commit() + + +def update_chat_greet(chat: Chat | ChatData, greeting: str): + db_chat = chat_dao.add_chat(chat) + db_chat.greet = greeting + commit() diff --git a/src/dao/quote.py b/kmua/dao/quote.py similarity index 85% rename from src/dao/quote.py rename to kmua/dao/quote.py index 2c1ef9b..faed69b 100644 --- a/src/dao/quote.py +++ b/kmua/dao/quote.py @@ -1,15 +1,15 @@ from telegram import Chat, Message, User -from ..models.models import ChatData, Quote, UserData -from .db import commit, db +from kmua.dao._db import commit, _db +from kmua.models import ChatData, Quote, UserData def get_quote_by_link(link: str) -> Quote | None: - return db.query(Quote).filter(Quote.link == link).first() + return _db.query(Quote).filter(Quote.link == link).first() def delete_quote(quote: Quote): - db.delete(quote) + _db.delete(quote) commit() @@ -43,7 +43,7 @@ def add_quote( link = f"https://t.me/c/{str(chat.id).removeprefix('-100')}/{message.id}" if quote := get_quote_by_link(link): return quote - db.add( + _db.add( Quote( chat_id=chat.id, user_id=user.id, @@ -59,4 +59,4 @@ def add_quote( def get_all_quotes_count() -> int: - return db.query(Quote).count() + return _db.query(Quote).count() diff --git a/src/dao/user.py b/kmua/dao/user.py similarity index 83% rename from src/dao/user.py rename to kmua/dao/user.py index 16edad1..5b688e6 100644 --- a/src/dao/user.py +++ b/kmua/dao/user.py @@ -1,11 +1,11 @@ from telegram import Chat, User -from ..models.models import ChatData, Quote, UserData -from .db import db, commit +from kmua.dao._db import commit, _db +from kmua.models import ChatData, Quote, UserData def get_user_by_id(user_id: int) -> UserData | None: - return db.query(UserData).filter(UserData.id == user_id).first() + return _db.query(UserData).filter(UserData.id == user_id).first() def add_user(user: User | UserData | Chat | ChatData) -> UserData: @@ -41,7 +41,7 @@ def add_user(user: User | UserData | Chat | ChatData) -> UserData: username=user.username, full_name=user.full_name, ) - db.add(userdata) + _db.add(userdata) commit() return get_user_by_id(user.id) @@ -51,14 +51,14 @@ def get_user_is_bot_global_admin(user: User | UserData) -> bool: def update_user_is_bot_global_admin(user: User | UserData, is_admin: bool): - db_user = add_user(user) - db_user.is_bot_global_admin = is_admin + _db_user = add_user(user) + _db_user.is_bot_global_admin = is_admin commit() def get_user_quotes(user: User | UserData) -> list[Quote]: - db_user = add_user(user) - return db_user.quotes + _db_user = add_user(user) + return _db_user.quotes def get_user_quotes_count(user: User | UserData) -> int: @@ -72,4 +72,4 @@ def get_user_quotes_page( def get_all_users_count() -> int: - return db.query(UserData).count() + return _db.query(UserData).count() diff --git a/src/service/user.py b/kmua/dao/user_service.py similarity index 50% rename from src/service/user.py rename to kmua/dao/user_service.py index bf2df11..2f0608a 100644 --- a/src/service/user.py +++ b/kmua/dao/user_service.py @@ -1,14 +1,13 @@ from telegram import Chat, User - -from ..dao.association import add_association_in_chat, get_association_in_chat_by_user -from ..dao.chat import get_chat_by_id -from ..dao import db -from ..dao.user import get_user_by_id -from ..models.models import ChatData, UserData +import kmua.dao.association as association_dao +import kmua.dao.chat as chat_dao +import kmua.dao.user as user_dao +from ._db import commit +from kmua.models import ChatData, UserData def get_user_is_bot_admin_in_chat(user: User | UserData, chat: Chat | ChatData) -> bool: - association = get_association_in_chat_by_user(chat, user) + association = association_dao.get_association_in_chat_by_user(chat, user) if association is None: return False return association.is_bot_admin @@ -17,20 +16,20 @@ def get_user_is_bot_admin_in_chat(user: User | UserData, chat: Chat | ChatData) def update_user_is_bot_admin_in_chat( user: User | UserData, chat: Chat | ChatData, is_admin: bool ): - association = get_association_in_chat_by_user(chat, user) + association = association_dao.get_association_in_chat_by_user(chat, user) if association is None: - add_association_in_chat(chat, user) - association = get_association_in_chat_by_user(chat, user) + association_dao.add_association_in_chat(chat, user) + association = association_dao.get_association_in_chat_by_user(chat, user) association.is_bot_admin = is_admin - db.commit() + commit() return association.is_bot_admin = is_admin - db.commit() + commit() def check_user_in_chat(user: User | UserData, chat: Chat | ChatData) -> bool: - db_user = get_user_by_id(user.id) - db_chat = get_chat_by_id(chat.id) + db_user = user_dao.get_user_by_id(user.id) + db_chat = chat_dao.get_chat_by_id(chat.id) if db_user is None or db_chat is None: return False return db_user in db_chat.members diff --git a/src/service/waifu.py b/kmua/dao/waifu.py similarity index 61% rename from src/service/waifu.py rename to kmua/dao/waifu.py index 4bbabd7..fb61a37 100644 --- a/src/service/waifu.py +++ b/kmua/dao/waifu.py @@ -2,37 +2,31 @@ from telegram import Chat, User -from ..dao.association import ( - get_association_in_chat_by_user, - get_associations_of_user, - get_associations_of_user_waifu_of, - get_associations_of_user_waifu_of_in_chat, - add_association_in_chat, - update_associations_all_waifu_id_to_none, -) -from ..dao.chat import add_chat, get_chat_by_id -from ..dao import db -from ..dao.user import add_user, get_user_by_id -from ..models.models import ChatData, UserData +import kmua.dao.association as association_dao +import kmua.dao.chat as chat_dao +import kmua.dao.user as user_dao +from ._db import commit + +from kmua.models import ChatData, UserData def _get_user_waifu_in_chat_common( user: User | UserData, chat: Chat | ChatData ) -> UserData | None: - db_user = get_user_by_id(user.id) + db_user = user_dao.get_user_by_id(user.id) if db_user is None: - add_user(user) + user_dao.add_user(user) return None - db_chat = get_chat_by_id(chat.id) + db_chat = chat_dao.get_chat_by_id(chat.id) if db_chat is None: - add_chat(chat) + chat_dao.add_chat(chat) return None - association = get_association_in_chat_by_user(chat, user) + association = association_dao.get_association_in_chat_by_user(chat, user) if association is None: return None if association.waifu_id is None: return None - return get_user_by_id(association.waifu_id) + return user_dao.get_user_by_id(association.waifu_id) def get_user_waifu_in_chat( @@ -40,9 +34,9 @@ def get_user_waifu_in_chat( ) -> UserData | None: waifu = _get_user_waifu_in_chat_common(user, chat) if waifu is None: - db_user = get_user_by_id(user.id) + db_user = user_dao.get_user_by_id(user.id) if db_user.married_waifu_id is not None: - return get_user_by_id(db_user.married_waifu_id) + return user_dao.get_user_by_id(db_user.married_waifu_id) return None return waifu @@ -63,24 +57,28 @@ def get_user_waifu_of_in_chat( :param chat: Chat or ChatData object :return: list of UserData object """ - db_user = get_user_by_id(user.id) + db_user = user_dao.get_user_by_id(user.id) if db_user is None: - add_user(user) + user_dao.add_user(user) return None - db_chat = get_chat_by_id(chat.id) + db_chat = chat_dao.get_chat_by_id(chat.id) if db_chat is None: - add_chat(chat) + chat_dao.add_chat(chat) return None - associations = get_associations_of_user_waifu_of_in_chat(db_user, db_chat) + associations = association_dao.get_associations_of_user_waifu_of_in_chat( + db_user, db_chat + ) if not associations: return None - return [get_user_by_id(association.user_id) for association in associations] + return [ + user_dao.get_user_by_id(association.user_id) for association in associations + ] def get_chat_married_users(chat: Chat | ChatData) -> list[UserData]: - db_chat = get_chat_by_id(chat.id) + db_chat = chat_dao.get_chat_by_id(chat.id) if db_chat is None: - add_chat(chat) + chat_dao.add_chat(chat) return [] return [user for user in db_chat.members if user.is_married] @@ -91,12 +89,12 @@ def get_chat_married_users_id(chat: Chat) -> list[int]: def get_user_married_waifu(user: User) -> UserData | None: - db_user = get_user_by_id(user.id) + db_user = user_dao.get_user_by_id(user.id) if db_user is None: - add_user(user) + user_dao.add_user(user) return None if married_waifu_id := db_user.married_waifu_id: - return get_user_by_id(married_waifu_id) + return user_dao.get_user_by_id(married_waifu_id) return None @@ -105,27 +103,27 @@ def put_user_waifu_in_chat( ) -> bool: if get_user_waifu_in_chat_exclude_married(user, chat) is not None: return False - if get_user_by_id(waifu.id) is None: - add_user(waifu) - if association := get_association_in_chat_by_user(chat, user): + if user_dao.get_user_by_id(waifu.id) is None: + user_dao.add_user(waifu) + if association := association_dao.get_association_in_chat_by_user(chat, user): if association.waifu_id is None: association.waifu_id = waifu.id - db.commit() + commit() return True else: return False else: - add_association_in_chat(chat, user, waifu) - db.commit() + association_dao.add_association_in_chat(chat, user, waifu) + commit() return True def refresh_user_waifu_in_chat(user: User | UserData, chat: Chat | ChatData): - association = get_association_in_chat_by_user(chat, user) + association = association_dao.get_association_in_chat_by_user(chat, user) if association is None: return association.waifu_id = None - db.commit() + commit() def get_chat_users_has_waifu(chat: Chat | ChatData) -> list[UserData]: @@ -135,7 +133,7 @@ def get_chat_users_has_waifu(chat: Chat | ChatData) -> list[UserData]: :param chat: Chat or ChatData object :return: list of UserData object """ - db_chat = add_chat(chat) + db_chat = chat_dao.add_chat(chat) return [ user for user in db_chat.members @@ -150,7 +148,7 @@ def get_chat_users_was_waifu(chat: Chat | ChatData) -> list[UserData]: :param chat: Chat or ChatData object :return: list of UserData object """ - db_chat = add_chat(chat) + db_chat = chat_dao.add_chat(chat) return [ user for user in db_chat.members @@ -171,41 +169,41 @@ def get_chat_user_participated_waifu(chat: Chat | ChatData) -> list[UserData]: def refresh_all_waifu_in_chat(chat: Chat | ChatData): - db_chat = add_chat(chat) + db_chat = chat_dao.add_chat(chat) for user in db_chat.members: refresh_user_waifu_in_chat(user, chat) async def refresh_all_waifu_data(): - update_associations_all_waifu_id_to_none() - db.commit() + association_dao.update_associations_all_waifu_id_to_none() + commit() def refresh_user_all_waifu(user: User | UserData): - db_user = add_user(user) - associations = get_associations_of_user(db_user) + db_user = user_dao.add_user(user) + associations = association_dao.get_associations_of_user(db_user) for association in associations: association.waifu_id = None - db.commit() + commit() def get_user_waifus(user: User | UserData) -> list[UserData]: - db_user = add_user(user) - associations = get_associations_of_user(db_user) + db_user = user_dao.add_user(user) + associations = association_dao.get_associations_of_user(db_user) return [ - get_user_by_id(association.waifu_id) + user_dao.get_user_by_id(association.waifu_id) for association in associations if association.waifu_id is not None ] def get_user_waifus_with_chat(user: User | UserData) -> list[tuple[UserData, ChatData]]: - db_user = add_user(user) - associations = get_associations_of_user(db_user) + db_user = user_dao.add_user(user) + associations = association_dao.get_associations_of_user(db_user) return [ ( - get_user_by_id(association.waifu_id), - get_chat_by_id(association.chat_id), + user_dao.get_user_by_id(association.waifu_id), + chat_dao.get_chat_by_id(association.chat_id), ) for association in associations if association.waifu_id is not None @@ -213,20 +211,22 @@ def get_user_waifus_with_chat(user: User | UserData) -> list[tuple[UserData, Cha def get_user_waifus_of(user: User | UserData) -> list[UserData]: - db_user = add_user(user) - associations = get_associations_of_user_waifu_of(db_user) - return [get_user_by_id(association.user_id) for association in associations] + db_user = user_dao.add_user(user) + associations = association_dao.get_associations_of_user_waifu_of(db_user) + return [ + user_dao.get_user_by_id(association.user_id) for association in associations + ] def get_user_waifus_of_with_chat( user: User | UserData, ) -> list[tuple[UserData, ChatData]]: - db_user = add_user(user) - associations = get_associations_of_user_waifu_of(db_user) + db_user = user_dao.add_user(user) + associations = association_dao.get_associations_of_user_waifu_of(db_user) return [ ( - get_user_by_id(association.user_id), - get_chat_by_id(association.chat_id), + user_dao.get_user_by_id(association.user_id), + chat_dao.get_chat_by_id(association.chat_id), ) for association in associations ] diff --git a/src/filters.py b/kmua/filters.py similarity index 100% rename from src/filters.py rename to kmua/filters.py diff --git a/src/handlers.py b/kmua/handlers.py similarity index 50% rename from src/handlers.py rename to kmua/handlers.py index f2d9cb7..0cf3288 100644 --- a/src/handlers.py +++ b/kmua/handlers.py @@ -4,161 +4,149 @@ CallbackQueryHandler, ChatMemberHandler, CommandHandler, - InlineQueryHandler, MessageHandler, filters, ) -from .callbacks.chatdata import chat_data_manage, chat_title_update -from .callbacks.chatinfo import getid -from .callbacks.chatmember import ( - on_member_join, - on_member_left, - set_greet, - track_chats, -) -from .callbacks.help import help -from .callbacks.keyword_reply import keyword_reply -from .callbacks.manage import ( - bot_data_refresh, - leave_chat, +from .callbacks import ( + chatdata, + chatinfo, + chatmember, + keyword_reply, manage, - refresh_waifu_data_manually, - set_bot_admin_globally, - set_bot_admin_in_chat, - status, -) -from .callbacks.others import chat_migration, error_notice_control -from .callbacks.quote import ( - delete_quote_in_chat, - inline_query_quote, + others, quote, - random_quote, - set_quote_probability, -) -from .callbacks.remake import remake -from .callbacks.slash import slash -from .callbacks.start import start -from .callbacks.sticker import sticker2img -from .callbacks.title import ( - set_title_permissions, - set_title_permissions_callback, + remake, + slash, + start, + sticker, + userdata, title, + waifu, + help, ) -from .callbacks.userdata import ( - delete_user_quote, - user_data_manage, - user_data_refresh, - user_waifu_manage, -) -from .callbacks.waifu import marry_waifu, remove_waifu, today_waifu, waifu_graph -from .config import settings from .filters import ( keyword_reply_filter, mention_or_private_filter, slash_filter, ) from .logger import logger +from .config import settings + # CommandHandlers -start_handler = CommandHandler("start", start, filters=mention_or_private_filter) +start_handler = CommandHandler("start", start.start, filters=mention_or_private_filter) -title_handler = CommandHandler("t", title, filters=filters.ChatType.GROUPS) +title_handler = CommandHandler("t", title.title, filters=filters.ChatType.GROUPS) -quote_handler = CommandHandler("q", quote, filters=filters.ChatType.GROUPS) +quote_handler = CommandHandler("q", quote.quote, filters=filters.ChatType.GROUPS) set_quote_probability_handler = CommandHandler( - "setqp", set_quote_probability, filters=filters.ChatType.GROUPS + "setqp", quote.set_quote_probability, filters=filters.ChatType.GROUPS ) delete_quote_handler = CommandHandler( - "d", delete_quote_in_chat, filters=filters.ChatType.GROUPS + "d", quote.delete_quote_in_chat, filters=filters.ChatType.GROUPS +) +qrand_handler = CommandHandler( + "qrand", quote.random_quote, filters=filters.ChatType.GROUPS ) -qrand_handler = CommandHandler("qrand", random_quote, filters=filters.ChatType.GROUPS) -help_handler = CommandHandler("help", help, filters=mention_or_private_filter) -error_notice_control_handler = CommandHandler("error_notice", error_notice_control) -remake_handler = CommandHandler("remake", remake) +help_handler = CommandHandler("help", help.help, filters=mention_or_private_filter) +error_notice_control_handler = CommandHandler( + "error_notice", others.error_notice_control +) +remake_handler = CommandHandler("remake", remake.remake) today_waifu_handler = CommandHandler( - "waifu", today_waifu, filters=filters.ChatType.GROUPS + "waifu", waifu.today_waifu, filters=filters.ChatType.GROUPS ) waifu_graph_handler = CommandHandler( - "waifu_graph", waifu_graph, filters=filters.ChatType.GROUPS + "waifu_graph", waifu.waifu_graph, filters=filters.ChatType.GROUPS ) set_greet_handler = CommandHandler( - "set_greet", set_greet, filters=filters.ChatType.GROUPS + "set_greet", chatmember.set_greet, filters=filters.ChatType.GROUPS ) -getid_handler = CommandHandler("id", getid) +getid_handler = CommandHandler("id", chatinfo.getid) set_title_permissions_handler = CommandHandler( - "sett", set_title_permissions, filters=filters.ChatType.GROUPS + "sett", title.set_title_permissions, filters=filters.ChatType.GROUPS ) chat_data_manage_handler = CommandHandler( - "manage", chat_data_manage, filters=filters.ChatType.GROUPS + "manage", chatdata.chat_data_manage, filters=filters.ChatType.GROUPS +) +bot_manage_handler = CommandHandler( + "manage", manage.manage, filters=filters.ChatType.PRIVATE ) -bot_manage_handler = CommandHandler("manage", manage, filters=filters.ChatType.PRIVATE) set_bot_admin_in_chat_handler = CommandHandler( - "set_bot_admin", set_bot_admin_in_chat, filters=filters.ChatType.GROUPS + "set_bot_admin", manage.set_bot_admin_in_chat, filters=filters.ChatType.GROUPS ) set_bot_admin_globally_handler = CommandHandler( - "set_bot_admin", set_bot_admin_globally, filters=filters.ChatType.PRIVATE + "set_bot_admin", manage.set_bot_admin_globally, filters=filters.ChatType.PRIVATE ) -leave_chat_handler = CommandHandler("leave_chat", leave_chat) +leave_chat_handler = CommandHandler("leave_chat", manage.leave_chat) refresh_waifu_data_manually_handler = CommandHandler( - "refresh_waifu_data", refresh_waifu_data_manually + "refresh_waifu_data", manage.refresh_waifu_data_manually ) -status_handler = CommandHandler("status", status) +status_handler = CommandHandler("status", manage.status) + # CallbackQueryHandlers -start_callback_handler = CallbackQueryHandler(start, pattern="back_home") +start_callback_handler = CallbackQueryHandler(start.start, pattern="back_home") -remove_waifu_handler = CallbackQueryHandler(remove_waifu, pattern="remove_waifu") +remove_waifu_handler = CallbackQueryHandler(waifu.remove_waifu, pattern="remove_waifu") user_waifu_manage_handler = CallbackQueryHandler( - user_waifu_manage, pattern="user_waifu_manage|set_waifu_mention|divorce" + userdata.user_waifu_manage, pattern="user_waifu_manage|set_waifu_mention|divorce" ) set_title_permissions_callback_handler = CallbackQueryHandler( - set_title_permissions_callback, pattern="set_title_permissions" + title.set_title_permissions_callback, pattern="set_title_permissions" ) chat_quote_manage_handler = CallbackQueryHandler( - delete_quote_in_chat, pattern="chat_quote_manage|delete_quote_in_chat" + quote.delete_quote_in_chat, pattern="chat_quote_manage|delete_quote_in_chat" ) user_data_manage_handler = CallbackQueryHandler( - user_data_manage, pattern="user_data_manage" + userdata.user_data_manage, pattern="user_data_manage" ) user_data_refresh_handler = CallbackQueryHandler( - user_data_refresh, pattern="user_data_refresh" + userdata.user_data_refresh, pattern="user_data_refresh" +) +marry_waifu_handler = CallbackQueryHandler( + waifu.marry_waifu, pattern=r".*marry_waifu.*" ) -marry_waifu_handler = CallbackQueryHandler(marry_waifu, pattern=r".*marry_waifu.*") user_quote_manage_handler = CallbackQueryHandler( - delete_user_quote, pattern="user_quote_manage|delete_user_quote" + userdata.delete_user_quote, pattern="user_quote_manage|delete_user_quote" ) bot_data_refresh_handler = CallbackQueryHandler( - bot_data_refresh, pattern="bot_data_refresh" + manage.bot_data_refresh, pattern="bot_data_refresh" ) -status_refresh_handler = CallbackQueryHandler(status, pattern="status_refresh") +status_refresh_handler = CallbackQueryHandler(manage.status, pattern="status_refresh") # others -chat_migration_handler = MessageHandler(filters.StatusUpdate.MIGRATE, chat_migration) -slash_handler = MessageHandler(slash_filter, slash) -inline_query_handler = InlineQueryHandler(inline_query_quote) +chat_migration_handler = MessageHandler( + filters.StatusUpdate.MIGRATE, others.chat_migration +) +slash_handler = MessageHandler(slash_filter, slash.slash) random_quote_handler = MessageHandler( - (~filters.COMMAND & filters.ChatType.GROUPS), random_quote + (~filters.COMMAND & filters.ChatType.GROUPS), quote.random_quote +) +keyword_reply_handler = MessageHandler( + keyword_reply_filter, keyword_reply.keyword_reply +) +track_chats_handler = ChatMemberHandler( + chatmember.track_chats, ChatMemberHandler.MY_CHAT_MEMBER ) -keyword_reply_handler = MessageHandler(keyword_reply_filter, keyword_reply) -track_chats_handler = ChatMemberHandler(track_chats, ChatMemberHandler.MY_CHAT_MEMBER) member_left_handler = MessageHandler( - filters.StatusUpdate.LEFT_CHAT_MEMBER, on_member_left + filters.StatusUpdate.LEFT_CHAT_MEMBER, chatmember.on_member_left ) member_join_handler = MessageHandler( - filters.StatusUpdate.NEW_CHAT_MEMBERS, on_member_join + filters.StatusUpdate.NEW_CHAT_MEMBERS, chatmember.on_member_join ) sticker2img_handler = MessageHandler( - (filters.Sticker.ALL & filters.ChatType.PRIVATE), sticker2img + (filters.Sticker.ALL & filters.ChatType.PRIVATE), sticker.sticker2img ) chat_title_update_handler = MessageHandler( filters.ChatType.GROUPS & filters.StatusUpdate.NEW_CHAT_TITLE, - chat_title_update, + chatdata.chat_title_update, ) handlers = [ @@ -200,7 +188,6 @@ start_callback_handler, remove_waifu_handler, slash_handler, - inline_query_handler, user_waifu_manage_handler, bot_data_refresh_handler, status_refresh_handler, diff --git a/src/logger.py b/kmua/logger.py similarity index 88% rename from src/logger.py rename to kmua/logger.py index 5ab8e06..02c8299 100644 --- a/src/logger.py +++ b/kmua/logger.py @@ -2,7 +2,7 @@ from loguru import logger -from .config import settings +from kmua.config import settings logger.remove() logger.add( diff --git a/src/models/models.py b/kmua/models.py similarity index 93% rename from src/models/models.py rename to kmua/models.py index 9d27d55..5c5463f 100644 --- a/src/models/models.py +++ b/kmua/models.py @@ -11,8 +11,11 @@ func, ) from sqlalchemy.orm import relationship +from sqlalchemy.ext.declarative import declarative_base +from kmua.dao._db import engine +from kmua.config import data_dir -from ..dao.db import Base +Base = declarative_base() class UserChatAssociation(Base): @@ -96,3 +99,9 @@ class Quote(Base): user = relationship("UserData", back_populates="quotes") chat = relationship("ChatData", back_populates="quotes") + + +if not data_dir.exists(): + data_dir.mkdir() + +Base.metadata.create_all(bind=engine) diff --git a/resource/TsukuA.ttc b/kmua/resource/TsukuA.ttc similarity index 100% rename from resource/TsukuA.ttc rename to kmua/resource/TsukuA.ttc diff --git a/resource/base.png b/kmua/resource/base.png similarity index 100% rename from resource/base.png rename to kmua/resource/base.png diff --git a/resource/word_dicts/weni.json b/kmua/resource/word_dicts/weni.json similarity index 100% rename from resource/word_dicts/weni.json rename to kmua/resource/word_dicts/weni.json diff --git a/pyproject.toml b/pyproject.toml index a0131bc..d7a1aaf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] -name = "kmua-bot" -version = "2.0.0" +name = "kmua" +version = "2.1.0" description = "" authors = ["krau "] readme = "README.md" diff --git a/src/dao/__init__.py b/src/dao/__init__.py deleted file mode 100644 index 1c75992..0000000 --- a/src/dao/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from .db import Base, engine -from ..models.models import ( - ChatData, # noqa: F401 - Quote, # noqa: F401 - UserData, # noqa: F401 - UserChatAssociation, # noqa: F401 -) - -from ..config import data_dir - -if not data_dir.exists(): - data_dir.mkdir() -Base.metadata.create_all(bind=engine) diff --git a/src/service/chat.py b/src/service/chat.py deleted file mode 100644 index ce174f4..0000000 --- a/src/service/chat.py +++ /dev/null @@ -1,35 +0,0 @@ -from telegram import Chat - -from ..dao.chat import get_chat_by_id, delete_chat -from ..dao import db -from ..dao.quote import delete_quote -from ..models.models import ChatData - - -def delete_chat_data_and_quotes(chat: Chat | ChatData): - # 删除与该群组相关的所有数据 - db_chat = get_chat_by_id(chat.id) - if db_chat is None: - return - for quote in db_chat.quotes: - delete_quote(quote) - delete_chat(db_chat) - db.commit() - - -def update_chat_id(old_id: int, new_id: int): - db_chat = get_chat_by_id(old_id) - if db_chat is None: - return - if get_chat_by_id(new_id) is not None: - return - db_chat.id = new_id - db.commit() - - -def update_chat_title(chat: Chat | ChatData, title: str): - db_chat = get_chat_by_id(chat.id) - if db_chat is None: - return - db_chat.title = title - db.commit()