From c3305f0f25a11eac01e4870bef72524e2104f88e Mon Sep 17 00:00:00 2001 From: "A.Shpak" Date: Tue, 12 Mar 2024 22:54:56 +0300 Subject: [PATCH] - --- chatushka/__init__.py | 7 ++-- chatushka/_chatushka.py | 56 +++++++++++++++++++++++++++-- chatushka/_matchers.py | 31 +++++++++++++++- chatushka/_models.py | 11 ++++-- chatushka/_transport.py | 18 ++++++++++ example/__main__.py | 78 +++++++++++++++++++++++++++++++---------- 6 files changed, 174 insertions(+), 27 deletions(-) diff --git a/chatushka/__init__.py b/chatushka/__init__.py index 58b067b..b4d9165 100644 --- a/chatushka/__init__.py +++ b/chatushka/__init__.py @@ -1,7 +1,7 @@ from chatushka._chatushka import ChatushkaBot from chatushka._errors import ChatushkaError, ChatushkaResponseError -from chatushka._matchers import BaseMatcher, CommandMatcher, RegExMatcher -from chatushka._models import Chat, Message, Update, User +from chatushka._matchers import BaseMatcher, CommandMatcher, RegExMatcher, EventMatcher +from chatushka._models import Chat, Message, Update, User, Events, ChatPermissions from chatushka._transport import TelegramBotAPI __all__ = [ @@ -12,9 +12,12 @@ # matchers "BaseMatcher", "CommandMatcher", + "EventMatcher", "RegExMatcher", # models + "Events", "Chat", + "ChatPermissions", "Message", "Update", "User", diff --git a/chatushka/_chatushka.py b/chatushka/_chatushka.py index 6e079c7..bd5b100 100644 --- a/chatushka/_chatushka.py +++ b/chatushka/_chatushka.py @@ -1,14 +1,23 @@ from asyncio import gather from collections.abc import Callable, Sequence -from typing import final +from contextlib import asynccontextmanager +from typing import final, MutableMapping from chatushka._constants import ( HTTP_POOLING_TIMEOUT, ) -from chatushka._matchers import CommandMatcher, Matcher, RegExMatcher +from chatushka._matchers import CommandMatcher, Matcher, RegExMatcher, EventMatcher +from chatushka._models import Events from chatushka._transport import TelegramBotAPI +@asynccontextmanager +async def _default_lifespan( + _: "ChatushkaBot", +) -> None: + yield + + @final class ChatushkaBot: def __init__( @@ -16,7 +25,10 @@ def __init__( *, token: str, cmd_prefixes: str | Sequence[str] = (), + lifespan=None, ) -> None: + self._state: MutableMapping = {} + self._lifespan = lifespan or _default_lifespan self._token = token if isinstance(cmd_prefixes, str): cmd_prefixes = [cmd_prefixes] @@ -97,6 +109,36 @@ def _wrapper( return _wrapper + def add_event( + self, + event: Events, + action: Callable, + chance_rate: float = 1.0, + ) -> None: + self.add_matcher( + EventMatcher( + event=event, + action=action, + chance_rate=chance_rate, + ) + ) + + def event( + self, + event: Events, + chance_rate: float = 1.0, + ) -> Callable: + def _wrapper( + func, + ) -> None: + self.add_event( + event=event, + action=func, + chance_rate=chance_rate, + ) + + return _wrapper + async def _check_updates( self, api: TelegramBotAPI, @@ -117,7 +159,7 @@ async def _check_updates( ) return offset - async def run( + async def _loop( self, ) -> None: offset: int | None = None @@ -130,3 +172,11 @@ async def run( api=api, offset=offset, ) + + async def run( + self, + ) -> None: + async with self._lifespan( + self, + ): + await self._loop() diff --git a/chatushka/_matchers.py b/chatushka/_matchers.py index ef412b2..6ac0ef6 100644 --- a/chatushka/_matchers.py +++ b/chatushka/_matchers.py @@ -5,7 +5,7 @@ from re import Pattern, compile from typing import TypeVar -from chatushka._models import Update +from chatushka._models import Update, Events from chatushka._transport import TelegramBotAPI Matcher = TypeVar("Matcher", bound="BaseMatcher") @@ -144,3 +144,32 @@ def _check( if not update.message or not update.message.text: return False return any(pattern.findall(update.message.text) for pattern in self._patterns) + + +class EventMatcher( + BaseMatcher, + ABC, +): + def __init__( + self, + event: Events, + action: Callable, + chance_rate: float = 1.0, + ) -> None: + super().__init__( + action=action, + chance_rate=chance_rate, + ) + self._event = event + + def _check( + self, + update: Update, + ) -> bool: + if update.message and update.message.text and self._event == "on_message": + return True + if update.message and update.message.new_chat_members and self._event == "on_new_chat_members": + return True + if update.message and update.message.new_chat_members and self._event == "on_left_chat_member": + return True + return False diff --git a/chatushka/_models.py b/chatushka/_models.py index 6f2592e..f241e93 100644 --- a/chatushka/_models.py +++ b/chatushka/_models.py @@ -1,6 +1,6 @@ from datetime import datetime -from enum import Enum -from typing import Optional +from enum import Enum, auto +from typing import Optional, Literal from pydantic import AliasChoices, BaseModel, ConfigDict, Field @@ -133,3 +133,10 @@ class ChatPermissions( can_send_media_messages: bool can_send_polls: bool can_send_other_messages: bool + + +Events = Literal[ + "on_message", + "on_new_chat_members", + "on_left_chat_member", +] diff --git a/chatushka/_transport.py b/chatushka/_transport.py index e35ad5d..747f9da 100644 --- a/chatushka/_transport.py +++ b/chatushka/_transport.py @@ -1,3 +1,4 @@ +from datetime import datetime from types import TracebackType from typing import Any, Literal @@ -11,6 +12,7 @@ ChatMemberStatuses, Message, Update, + ChatPermissions, ) @@ -131,3 +133,19 @@ async def get_chat_administrators( if status == ChatMemberStatuses.ADMINISTRATOR: admins.append(ChatMemberAdministrator.model_validate(result)) return admins + + async def restrict_chat_member( + self, + chat_id: int, + user_id: int, + permissions: ChatPermissions, + until_date: datetime, + ) -> bool: + result = await self._api_request( + "restrictChatMember", + chat_id=chat_id, + user_id=user_id, + permissions=permissions.json(), + until_date=int(until_date.timestamp()), + ) + return result diff --git a/example/__main__.py b/example/__main__.py index 3ec0e92..4302eef 100644 --- a/example/__main__.py +++ b/example/__main__.py @@ -1,5 +1,5 @@ import asyncio -import contextlib +import datetime import os import random @@ -52,27 +52,17 @@ async def luk_handler( async def id_handler( api: chatushka.TelegramBotAPI, message: chatushka.Message, + user: chatushka.User, + chat: chatushka.Chat, ) -> None: - admins = [] - with contextlib.suppress(chatushka.ChatushkaResponseError): - admins = await api.get_chat_administrators( - chat_id=message.chat.id, - ) - line_tmpl = "{id_type}: 
{id_value}
" - ids = {"user_id": message.user.id} - for admin in admins: - if admin.user.id == message.user.id: - ids = ids | {"chat_id": message.chat.id} - text = "\n".join( - [ - line_tmpl.format(id_type=id_type, id_value=id_value) - for id_type, id_value in ids.items() - ] - ) + text = f"*USER_ID:* `{user.id}`" + if user.id != chat.id: + text += f"\n*CHAT_ID:* `{chat.id}`" await api.send_message( - chat_id=message.chat.id, + chat_id=chat.id, text=text, reply_to_message_id=message.message_id, + parse_mode="markdown", ) @@ -88,5 +78,55 @@ async def ping_handler( ) +@bot.event("on_message") +async def welcome_handler( + api: chatushka.TelegramBotAPI, + chat: chatushka.Chat, + message: chatushka.Message, +) -> None: + await api.send_message( + chat_id=chat.id, + text="message", + reply_to_message_id=message.message_id, + ) + + +@bot.cmd("suicide", "wtf") +async def suicide_handler( + api: chatushka.TelegramBotAPI, + message: chatushka.Message, +) -> None: + restrict_time = datetime.timedelta(minutes=random.randrange(1, 4 * 60)) + try: + is_success = await api.restrict_chat_member( + chat_id=message.chat.id, + user_id=message.user.id, + permissions=chatushka.ChatPermissions( + can_send_messages=False, + can_send_media_messages=False, + can_send_polls=False, + can_send_other_messages=False, + ), + until_date=datetime.datetime.now( + tz=datetime.timezone.utc, + ) + restrict_time, + ) + except ValueError: + is_success = False + if is_success: + await api.send_message( + chat_id=message.chat.id, + text=f"Пользователь {message.user.readable_name} самовыпилился на {restrict_time}", + ) + return None + await api.send_message( + chat_id=message.chat.id, + text=f"Лапки коротковаты чтоб убить {message.user.readable_name}", + reply_to_message_id=message.message_id, + ) + + if __name__ == "__main__": - asyncio.run(bot.run()) + asyncio.run( + bot.run(), + )