From 0de7969d30e3e57d681afdfcadd245f6988a0342 Mon Sep 17 00:00:00 2001 From: Igor Tolmachov Date: Mon, 5 Dec 2022 22:29:41 +0900 Subject: 2.1 --- handlers/config.py | 23 +++++++++++++++----- handlers/gen.py | 59 ++++++-------------------------------------------- handlers/member.py | 23 +++++++++++++++++--- handlers/middleware.py | 5 +++-- handlers/pin.py | 6 ++--- poetry.lock | 41 ++++++++++++++++++++++++----------- pyproject.toml | 2 +- requirements.txt | 3 ++- shared/__init__.py | 1 + shared/commands.py | 3 --- shared/database.py | 1 + shared/instances.py | 6 ++--- shared/samples.py | 31 ++++++++++++++++++++++++++ shared/settings.py | 49 ++++++++++++++++++++--------------------- test.py | 15 +++++++++++++ utils/filters.py | 11 ++++++++-- 16 files changed, 164 insertions(+), 115 deletions(-) create mode 100644 shared/samples.py create mode 100644 test.py diff --git a/handlers/config.py b/handlers/config.py index 7628805..2865694 100644 --- a/handlers/config.py +++ b/handlers/config.py @@ -1,3 +1,4 @@ +from copy import deepcopy from json import JSONDecodeError, dumps, loads from logging import info @@ -5,7 +6,8 @@ from aiogram import types as t from pydantic import BaseModel, ValidationError from shared.database import Message -from shared.instances import config, dp, session +from shared.instances import chats, dp, session +from shared.samples import samples from shared.settings import Config from utils import filters as f @@ -13,9 +15,19 @@ from utils import filters as f @dp.message_handler(f.user.is_admin, commands=["void"]) async def void_command(msg: t.Message) -> None: if msg.get_args() == "Я знаю что делаю": + samples.delete(msg.chat.id) with session.begin() as s: - s.query(Message).filter(Message.chat_id == msg.chat.id).delete() - await msg.answer("Лоботомия проведена успешно") + query = s.query(Message).filter(Message.chat_id == msg.chat.id) + if msg.reply_to_message is not None: + query = query.filter(Message.user_id == msg.reply_to_message.from_user.id) + query.delete() + if msg.reply_to_message is not None: + await msg.answer( + f'Связи пользователя {msg.reply_to_message.from_user.mention} были очищены', + parse_mode=t.ParseMode.HTML, + ) + else: + await msg.answer("Связи были очищены") else: await msg.answer( "Напишите /void Я знаю что делаю", @@ -79,7 +91,7 @@ async def settings_command(msg: t.Message) -> None: return text - chat_config = config.get_config(msg.chat.id) + chat_config = deepcopy(chats.get(msg.chat.id)) args = msg.get_args().split() if len(args) == 0: text = f"/config{get_fields(chat_config)}" @@ -88,8 +100,7 @@ async def settings_command(msg: t.Message) -> None: elif len(args) == 2: try: text = set_filed(chat_config, args[0].split("."), args[1]) - config.set_config(msg.chat.id, Config.parse_obj(chat_config.dict())) - config.save("data/config.json") + chats.set(msg.chat.id, Config.parse_obj(chat_config.dict())) except (JSONDecodeError, ValidationError): text = "Неверное значение" else: diff --git a/handlers/gen.py b/handlers/gen.py index 72afb79..ef397db 100644 --- a/handlers/gen.py +++ b/handlers/gen.py @@ -1,60 +1,15 @@ -import random - -import mc from aiogram import types as t -from shared.database import Message -from shared.instances import bot, config, dp, session +from shared.instances import chats, dp +from shared.samples import samples from utils import filters as f -def get_text(chat_id: int) -> str: - with session() as s: - samples = [ - m.tuple()[0] - for m in s.query(Message.message).filter(Message.chat_id == chat_id).all() - ] - - assert ( - len(samples) != 0 - ), "Нету данных на основе которых можно сгенерировать сообщение" - - generator = mc.PhraseGenerator(samples) - gen_config = config.get_config(chat_id).gen - validators = [] - - if gen_config.max_word_count is not None or gen_config.min_word_count is not None: - validators.append( - mc.builtin.validators.words_count( - minimal=gen_config.min_word_count, - maximal=gen_config.max_word_count, - ) - ) - - while True: - message = generator.generate_phrase_or_none(1, validators=validators) - if message is not None: - return message - - @dp.message_handler(commands=["gen"]) async def gen_command(msg: t.Message) -> None: - if config.get_config(msg.chat.id).gen.delete_command: + if chats.get(msg.chat.id).gen.delete_command: await msg.delete() - await msg.answer(get_text(msg.chat.id)) - - -@dp.message_handler(commands=["del"]) -async def del_command(msg: t.Message) -> None: - await msg.delete() - - if msg.reply_to_message: - if msg.reply_to_message.from_user.id == bot.id: - await msg.reply_to_message.delete() - else: - await msg.reply("Можно удалять только сообщения бота") - else: - await msg.reply("Вы не ответили на сообщение") + await msg.answer(samples.get(msg.chat.id).generate()) @dp.message_handler( @@ -63,7 +18,7 @@ async def del_command(msg: t.Message) -> None: content_types=[t.ContentType.ANY], ) async def chance_message(msg: t.Message) -> None: - if config.get_config(msg.chat.id).gen.reply: - await msg.reply(get_text(msg.chat.id)) + if chats.get(msg.chat.id).gen.reply: + await msg.reply(samples.get(msg.chat.id).generate()) else: - await msg.answer(get_text(msg.chat.id)) + await msg.answer(samples.get(msg.chat.id).generate()) diff --git a/handlers/member.py b/handlers/member.py index 830351d..b228f6c 100644 --- a/handlers/member.py +++ b/handlers/member.py @@ -1,17 +1,22 @@ from aiogram import types as t -from shared.instances import bot, config, dp +from shared.instances import bot, chats, dp from utils import filters as f +polls: dict[int, int] = {} + @dp.chat_join_request_handler() async def new_member(cjr: t.ChatJoinRequest) -> None: + if cjr.from_user.id in polls: + return + reply = await bot.send_message( cjr.chat.id, f'{cjr.from_user.mention} хочет в чат', parse_mode=t.ParseMode.HTML, ) - await reply.reply_poll( + poll = await reply.reply_poll( "Пускаем ?", [ "Да", @@ -28,6 +33,18 @@ async def new_member(cjr: t.ChatJoinRequest) -> None: await bot.send_message( cjr.from_user.id, "Заявка на вступление в группу будет вскоре рассмотрена" ) + polls[cjr.from_user.id] = poll.message_id + + +@dp.chat_member_handler(f.user.new_user) +async def admin_accept_member(cmu: t.ChatMemberUpdated) -> None: + user_id = cmu.new_chat_member.user.id + if user_id in polls: + polls.pop(user_id) + await bot.send_message( + user_id, + "Ваша заявка на вступление принята, добро пожаловать в группу", + ) @dp.callback_query_handler( @@ -38,7 +55,7 @@ async def check_poll(clb: t.CallbackQuery) -> None: msg = clb.message data = clb.data.split(":") user_id = int(data[1]) - min_answers = config.get_config(msg.chat.id).members.answer_count + min_answers = chats.get(msg.chat.id).members.answer_count if poll.total_voter_count < min_answers: await clb.answer( diff --git a/handlers/middleware.py b/handlers/middleware.py index 6cf39b6..046037d 100644 --- a/handlers/middleware.py +++ b/handlers/middleware.py @@ -1,10 +1,9 @@ -from logging import info - from aiogram import types as t from aiogram.dispatcher.middlewares import BaseMiddleware from shared.database import Message from shared.instances import session +from shared.samples import samples class MessageMiddleware(BaseMiddleware): @@ -16,6 +15,8 @@ class MessageMiddleware(BaseMiddleware): Message( chat_id=msg.chat.id, message_id=msg.message_id, + user_id=msg.from_user.id, message=text, ) ) + samples.get(msg.chat.id).feed(text) diff --git a/handlers/pin.py b/handlers/pin.py index 17995ee..e4d21db 100644 --- a/handlers/pin.py +++ b/handlers/pin.py @@ -1,6 +1,6 @@ from aiogram import types as t -from shared.instances import bot, config, dp +from shared.instances import bot, chats, dp from utils import filters as f @@ -18,7 +18,7 @@ async def pin_command(msg: t.Message) -> None: "Да", "Нет", ], - config.get_config(msg.chat.id).pin.anonym, + chats.get(msg.chat.id).pin.anonym, reply_markup=t.InlineKeyboardMarkup().add( t.InlineKeyboardButton( "Проверить опрос", @@ -37,7 +37,7 @@ async def check_poll(clb: t.CallbackQuery) -> None: poll = clb.message.poll msg = clb.message pin = int(clb.data.split(":")[1]) - min_answers = config.get_config(msg.chat.id).pin.answer_count + min_answers = chats.get(msg.chat.id).pin.answer_count if poll.total_voter_count < min_answers: await clb.answer( diff --git a/poetry.lock b/poetry.lock index 2ef1bd2..01570d0 100644 --- a/poetry.lock +++ b/poetry.lock @@ -189,14 +189,6 @@ category = "main" optional = false python-versions = ">=3.7,<4.0" -[[package]] -name = "mc.py" -version = "4.0.0" -description = "Python package which provides you a simple way to generate phrases using Markov chains" -category = "main" -optional = false -python-versions = ">=3.9,<4.0" - [[package]] name = "multidict" version = "6.0.2" @@ -310,6 +302,17 @@ postgresql_psycopg2cffi = ["psycopg2cffi"] pymysql = ["pymysql"] sqlcipher = ["sqlcipher3-binary"] +[[package]] +name = "tgen" +version = "0.3.0" +description = "A pretty simple text generation library" +category = "main" +optional = false +python-versions = ">=3.10,<4.0" + +[package.dependencies] +u-msgpack-python = ">=2.7.2,<3.0.0" + [[package]] name = "tomli" version = "2.0.1" @@ -326,6 +329,14 @@ category = "main" optional = false python-versions = ">=3.7" +[[package]] +name = "u-msgpack-python" +version = "2.7.2" +description = "A portable, lightweight MessagePack serializer and deserializer written in pure Python." +category = "main" +optional = false +python-versions = "*" + [[package]] name = "yarl" version = "1.8.1" @@ -341,7 +352,7 @@ multidict = ">=4.0" [metadata] lock-version = "1.1" python-versions = "^3.10" -content-hash = "4a5eec948c56acc838b2883ca7324b305189367039e1974947768f4ac252db6b" +content-hash = "a60189eae7993f31bd8274d5f3710cde0e406b80ac5fc7813e0fbf02ea94331e" [metadata.files] aiogram = [ @@ -642,10 +653,6 @@ magic-filter = [ {file = "magic-filter-1.0.9.tar.gz", hash = "sha256:d0f1ffa5ff1fbe5105fd5f293c79b5d3795f336ea0f6129c636959a687bf422a"}, {file = "magic_filter-1.0.9-py3-none-any.whl", hash = "sha256:51002312a8972fa514b998b7ff89340c98be3fc499967c1f5f2af98d13baf8d5"}, ] -"mc.py" = [ - {file = "mc.py-4.0.0-py3-none-any.whl", hash = "sha256:aa9a13bf95fc818cc56d1363f3522b050f5ad92715218701992b23217b151417"}, - {file = "mc.py-4.0.0.tar.gz", hash = "sha256:310f2e9a48ca69b23d4e538a316575380654ff94a14f6e0e38b85928d5d5ea19"}, -] multidict = [ {file = "multidict-6.0.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:0b9e95a740109c6047602f4db4da9949e6c5945cefbad34a1299775ddc9a62e2"}, {file = "multidict-6.0.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ac0e27844758d7177989ce406acc6a83c16ed4524ebc363c1f748cba184d89d3"}, @@ -836,6 +843,10 @@ SQLAlchemy = [ {file = "SQLAlchemy-2.0.0b3-py3-none-any.whl", hash = "sha256:5a4dcd4e740cf0e1c38f4053b365a0270690dee60e6061f0eebb1b8d4afcc455"}, {file = "SQLAlchemy-2.0.0b3.tar.gz", hash = "sha256:8325a4648e639cb2010199f64fad679d2f4ec8ce7e6f424ee1a41b07940cadb6"}, ] +tgen = [ + {file = "tgen-0.3.0-py3-none-any.whl", hash = "sha256:6c145c1f99613cfa345251a5fee9333b648ae2b42af39a7f2d5ddd5449d88e29"}, + {file = "tgen-0.3.0.tar.gz", hash = "sha256:459d6ef8490d158fdfef90ea429fe44a0baf8b0d61fbd98c2694ddb183a526a8"}, +] tomli = [ {file = "tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc"}, {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, @@ -844,6 +855,10 @@ typing-extensions = [ {file = "typing_extensions-4.4.0-py3-none-any.whl", hash = "sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e"}, {file = "typing_extensions-4.4.0.tar.gz", hash = "sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa"}, ] +u-msgpack-python = [ + {file = "u-msgpack-python-2.7.2.tar.gz", hash = "sha256:e86f7ac6aa0ef4c6c49f004b4fd435bce99c23e2dd5d73003f3f9816024c2bd8"}, + {file = "u_msgpack_python-2.7.2-py2.py3-none-any.whl", hash = "sha256:00c77bbb65f8f68c8bd2570e0a14dee84aba429629ea78adc713f94da52ee386"}, +] yarl = [ {file = "yarl-1.8.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:abc06b97407868ef38f3d172762f4069323de52f2b70d133d096a48d72215d28"}, {file = "yarl-1.8.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:07b21e274de4c637f3e3b7104694e53260b5fc10d51fb3ec5fed1da8e0f754e3"}, diff --git a/pyproject.toml b/pyproject.toml index fe03f00..3faa60d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,8 +9,8 @@ license = "GPL-3.0-only" python = "^3.10" aiogram = "^2.20" SQLAlchemy = "^2b3" -"mc.py" = "^4.0.0" pydantic = "^1.10.2" +tgen = "^0.3.0" [tool.poetry.group.dev.dependencies] mypy = "^0.991" diff --git a/requirements.txt b/requirements.txt index 90d4534..cf90e66 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,10 +10,11 @@ frozenlist==1.3.3 ; python_version >= "3.10" and python_version < "4.0" greenlet==2.0.1 ; python_version >= "3.10" and python_version < "4.0" and platform_machine == "aarch64" or python_version >= "3.10" and python_version < "4.0" and platform_machine == "ppc64le" or python_version >= "3.10" and python_version < "4.0" and platform_machine == "x86_64" or python_version >= "3.10" and python_version < "4.0" and platform_machine == "amd64" or python_version >= "3.10" and python_version < "4.0" and platform_machine == "AMD64" or python_version >= "3.10" and python_version < "4.0" and platform_machine == "win32" or python_version >= "3.10" and python_version < "4.0" and platform_machine == "WIN32" idna==3.4 ; python_version >= "3.10" and python_version < "4.0" magic-filter==1.0.9 ; python_version >= "3.10" and python_version < "4.0" -mc-py==4.0.0 ; python_version >= "3.10" and python_version < "4.0" multidict==6.0.2 ; python_version >= "3.10" and python_version < "4.0" pydantic==1.10.2 ; python_version >= "3.10" and python_version < "4.0" pytz==2022.6 ; python_version >= "3.10" and python_version < "4.0" sqlalchemy==2.0.0b3 ; python_version >= "3.10" and python_version < "4.0" +tgen==0.3.0 ; python_version >= "3.10" and python_version < "4.0" typing-extensions==4.4.0 ; python_version >= "3.10" and python_version < "4.0" +u-msgpack-python==2.7.2 ; python_version >= "3.10" and python_version < "4.0" yarl==1.8.1 ; python_version >= "3.10" and python_version < "4.0" diff --git a/shared/__init__.py b/shared/__init__.py index f708bfe..edfd5c2 100644 --- a/shared/__init__.py +++ b/shared/__init__.py @@ -2,4 +2,5 @@ from . import settings from . import instances from . import database +from . import samples from . import commands diff --git a/shared/commands.py b/shared/commands.py index 684997c..a7ff002 100644 --- a/shared/commands.py +++ b/shared/commands.py @@ -6,19 +6,16 @@ from aiogram.types import BotCommandScopeAllPrivateChats as private commands = { group(): [ cmd("gen", "Сгенерировать сообщение"), - cmd("del", "Удалить сообщение бота"), cmd("pin", "Создать опрос для закрепления сообщения"), ], admin(): [ cmd("gen", "Сгенерировать сообщение"), - cmd("del", "Удалить сообщение бота"), cmd("pin", "Создать опрос для закрепления сообщения"), cmd("void", "Отчистить связи для генерации сообщений"), cmd("config", "Открыть настройки чата"), ], private(): [ cmd("gen", "Сгенерировать сообщение"), - cmd("del", "Удалить сообщение бота"), cmd("void", "Отчистить связи для генерации сообщений"), ], } diff --git a/shared/database.py b/shared/database.py index 9dfe868..6ce5fa2 100644 --- a/shared/database.py +++ b/shared/database.py @@ -11,6 +11,7 @@ class Message(Base): __tablename__ = "messages" chat_id: Mapped[int] = mapped_column(primary_key=True) message_id: Mapped[int] = mapped_column(primary_key=True) + user_id: Mapped[int] message: Mapped[str] diff --git a/shared/instances.py b/shared/instances.py index fb511b3..5528206 100644 --- a/shared/instances.py +++ b/shared/instances.py @@ -7,10 +7,10 @@ from sqlalchemy.orm import Session, sessionmaker from shared.settings import Chats, Settings settings = Settings() -config = Chats() +chats = Chats("data/config.json") if not exists("data/config.json"): - config.save("data/config.json") -config.load("data/config.json") + chats.save() +chats.load() bot = Bot(token=settings.token) dp = Dispatcher(bot) diff --git a/shared/samples.py b/shared/samples.py new file mode 100644 index 0000000..740cff5 --- /dev/null +++ b/shared/samples.py @@ -0,0 +1,31 @@ +from tgen import TextGenerator + +from shared.database import Message +from shared.instances import session + + +class Samples: + samples: dict[int, TextGenerator] + + def __init__(self) -> None: + self.samples = {} + + def get(self, chat_id: int) -> TextGenerator: + if chat_id not in self.samples: + with session() as s: + samples = [ + m.tuple()[0] + for m in s.query(Message.message) + .filter(Message.chat_id == chat_id) + .all() + ] + + self.samples[chat_id] = TextGenerator.from_samples(samples) + return self.samples[chat_id] + + def delete(self, chat_id: int) -> None: + if chat_id in self.samples: + self.samples.pop(chat_id) + + +samples = Samples() diff --git a/shared/settings.py b/shared/settings.py index 560f8f8..bff7ff9 100644 --- a/shared/settings.py +++ b/shared/settings.py @@ -11,7 +11,7 @@ class GenConfig(BaseModel): chance: int = Field( 10, description="Шанс с которым бот сгенерирует сообщение", - ge=1, + ge=0, le=100, ) reply: bool = Field( @@ -22,16 +22,6 @@ class GenConfig(BaseModel): True, description="Включить/Выключить удаление /gen команды", ) - min_word_count: int | str | None = Field( - None, - description="Минимальное количество слов в сгенерированном предложении", - ge=1, - ) - max_word_count: int | None = Field( - None, - description="Максимальное количество слов в сгенерированном предложении", - ge=1, - ) class PollConfig(BaseModel): @@ -60,22 +50,29 @@ class Config(BaseModel): ) -class Chats(BaseModel): - chats: dict[int, Config] = {} +class Chats: + file_name: str + configs: dict[int, Config] + + def __init__(self, file_name: str) -> None: + self.file_name = file_name + self.configs = {} - @classmethod - def load(cls, file_name: str) -> "Chats": - with open(file_name, "r") as file: - return cls.parse_obj(load(file)) + def load(self) -> None: + with open(self.file_name, "r") as file: + self.configs = { + id_: Config.parse_obj(config) for id_, config in load(file).items() + } - def save(self, file_name: str) -> None: - with open(file_name, "w") as file: - dump(self.schema(), file) + def save(self) -> None: + with open(self.file_name, "w") as file: + dump({id_: config.dict() for id_, config in self.configs.items()}, file) - def get_config(self, chat_id: int) -> Config: - if chat_id not in self.chats: - self.chats[chat_id] = Config() - return self.chats[chat_id] + def get(self, chat_id: int) -> Config: + if chat_id not in self.configs: + self.configs[chat_id] = Config() + return self.configs[chat_id] - def set_config(self, chat_id: int, config: Config) -> None: - self.chats[chat_id] = config + def set(self, chat_id: int, config: Config) -> None: + self.configs[chat_id] = config + self.save() diff --git a/test.py b/test.py new file mode 100644 index 0000000..f6fe3a0 --- /dev/null +++ b/test.py @@ -0,0 +1,15 @@ +class A: + def __init__(self, b: dict = {}) -> None: + self.b = b + + +a1 = A() +a2 = A() + +print(a1.b) +print(a2.b) + +a1.b["test"] = 123 + +print(a1.b) +print(a2.b) diff --git a/utils/filters.py b/utils/filters.py index 59302fb..4f0fee2 100644 --- a/utils/filters.py +++ b/utils/filters.py @@ -5,7 +5,7 @@ from random import randint from aiogram import filters as f from aiogram import types as t -from shared.instances import config +from shared.instances import chats class message: @@ -13,7 +13,7 @@ class message: @staticmethod def chance(msg: t.Message) -> bool: - return config.get_config(msg.chat.id).gen.chance >= randint(1, 100) + return chats.get(msg.chat.id).gen.chance >= randint(1, 100) class user: @@ -28,3 +28,10 @@ class user: ) return False return True + + @staticmethod + async def new_user(cmu: t.ChatMemberUpdated) -> bool: + return ( + not cmu.old_chat_member.is_chat_member() + and cmu.new_chat_member.is_chat_member() + ) -- cgit v1.2.3