Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve Dispatcher #39

Merged
merged 8 commits into from
Aug 3, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion tgram/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
__all__ = ["types", "TgBot", "handlers", "filters"]
__all__ = ["types", "TgBot", "handlers", "filters", "compose", "StopPropagation"]
__version__ = "1.10.7"

__author__ = "Zaid"
@@ -7,3 +7,6 @@

from . import types, handlers, filters
from .tgbot import TgBot
from .sync import compose

from .errors import StopPropagation
4 changes: 4 additions & 0 deletions tgram/errors.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,6 @@
class APIException(Exception):
pass


class StopPropagation(Exception):
pass
6 changes: 6 additions & 0 deletions tgram/sync.py
Original file line number Diff line number Diff line change
@@ -4,6 +4,8 @@
import threading
import logging

from tgram import utils

logger = logging.getLogger(__name__)


@@ -107,3 +109,7 @@ def wrap(source):
method
):
async_to_sync(source, name)


async_to_sync(utils, "compose")
compose = utils.compose
33 changes: 19 additions & 14 deletions tgram/tgbot.py
Original file line number Diff line number Diff line change
@@ -29,10 +29,6 @@


class Dispatcher:
_is_running = False
_handlers: List["tgram.handlers.Handler"] = []
_listen_handlers: List["tgram.types.Listener"] = []

async def run_for_updates(self: "TgBot", skip_updates: bool = True) -> None:
if self.plugins:
self.load_plugins()
@@ -41,10 +37,10 @@ async def run_for_updates(self: "TgBot", skip_updates: bool = True) -> None:
self.allowed_updates,
100,
)
self._is_running = True
self.is_running = True
self.me = await self.get_me()

while self._is_running:
while self.is_running:
try:
updates = await self.get_updates(
offset=offset,
@@ -56,13 +52,19 @@ async def run_for_updates(self: "TgBot", skip_updates: bool = True) -> None:
offset = update.update_id + 1
await self._check_update(update)
except (asyncio.CancelledError, KeyboardInterrupt):
self._is_running = False
self.is_running = False
except tgram.StopPropagation:
pass
except Exception as e:
logger.exception(e)

session = await self._get_session()
await session.close()

async def stop(self) -> Literal[True]:
self.is_running = False
return True

async def _check_cancel(self: "TgBot", callback: Callable, update: Any) -> bool:
logger.debug("Checking listener in %s func", callback.__name__)
try:
@@ -121,11 +123,6 @@ async def _process_update(self: "TgBot", update: Any, callback: Callable) -> Non


class TgBot(TelegramBotMethods, Decorators, Dispatcher):
me: "tgram.types.User" = None
_session: "aiohttp.ClientSession" = None
_api_url: str = None
_custom_types: dict = {}

def __init__(
self,
bot_token: str,
@@ -150,10 +147,18 @@ def __init__(
self.executor = ThreadPoolExecutor(self.workers, thread_name_prefix="Handlers")
self.loop = asyncio.get_event_loop()

self.is_running: bool = None
self.me: "tgram.types.User" = None

self._listen_handlers: List["tgram.types.Listener"] = []
self._handlers: List["tgram.handlers.Handler"] = []
self._custom_types: dict = {}
self._session: "aiohttp.ClientSession" = None

if not api_url.endswith("/"):
api_url += "/"

self._api_url = f"{api_url}bot{bot_token}/"
self._api_url: str = f"{api_url}bot{bot_token}/"

def add_handler(self, handler: "tgram.handlers.Handler") -> None:
if handler.type == "all":
@@ -226,7 +231,7 @@ async def _send_request(self, method: str, **kwargs) -> Any:
),
)

if not self._is_running:
if not self.is_running:
await session.close()

response_json = await response.json()
7 changes: 7 additions & 0 deletions tgram/utils.py
Original file line number Diff line number Diff line change
@@ -2,6 +2,7 @@
import tgram
import re
import html
import asyncio

from pathlib import Path
from typing import List, Union
@@ -315,3 +316,9 @@ def recursive(entity_i: int) -> int:
last_offset = offset

return remove_surrogates(text)


async def compose(bots: List["tgram.TgBot"]):
tasks = [asyncio.create_task(bot.run_for_updates()) for bot in bots]

return await asyncio.wait(tasks)