diff --git a/aiogram/bot/base.py b/aiogram/bot/base.py index 19031fd0..4f55dbd7 100644 --- a/aiogram/bot/base.py +++ b/aiogram/bot/base.py @@ -1,11 +1,14 @@ import asyncio +import contextlib import io import ssl import typing +from contextvars import ContextVar from typing import Dict, List, Optional, Union import aiohttp import certifi +from aiohttp.helpers import sentinel from . import api from ..types import ParseMode, base @@ -17,6 +20,7 @@ class BaseBot: """ Base class for bot. It's raw bot. """ + _ctx_timeout = ContextVar('TelegramRequestTimeout') def __init__( self, @@ -27,6 +31,7 @@ class BaseBot: proxy_auth: Optional[aiohttp.BasicAuth] = None, validate_token: Optional[base.Boolean] = True, parse_mode: typing.Optional[base.String] = None, + timeout: typing.Optional[typing.Union[base.Integer, base.Float]] = None ): """ Instructions how to get Bot token is found here: https://core.telegram.org/bots#3-how-do-i-create-a-bot @@ -45,6 +50,8 @@ class BaseBot: :type validate_token: :obj:`bool` :param parse_mode: You can set default parse mode :type parse_mode: :obj:`str` + :param timeout: Request timeout + :type timeout: :obj:`typing.Optional[typing.Union[base.Integer, base.Float]]` :raise: when token is invalid throw an :obj:`aiogram.utils.exceptions.ValidationError` """ # Authentication @@ -83,11 +90,51 @@ class BaseBot: self.proxy_auth = None else: connector = aiohttp.TCPConnector(limit=connections_limit, ssl=ssl_context, loop=self.loop) + self._timeout = None + self.timeout = timeout self.session = aiohttp.ClientSession(connector=connector, loop=self.loop, json_serialize=json.dumps) self.parse_mode = parse_mode + @staticmethod + def _prepare_timeout( + value: typing.Optional[typing.Union[base.Integer, base.Float, aiohttp.ClientTimeout]] + ) -> typing.Optional[aiohttp.ClientTimeout]: + if value is None or isinstance(value, aiohttp.ClientTimeout): + return value + return aiohttp.ClientTimeout(total=value) + + @property + def timeout(self): + timeout = self._ctx_timeout.get(self._timeout) + if timeout is None: + return sentinel + return timeout + + @timeout.setter + def timeout(self, value): + self._timeout = self._prepare_timeout(value) + + @timeout.deleter + def timeout(self): + self.timeout = None + + @contextlib.contextmanager + def request_timeout(self, timeout): + """ + Context manager implements opportunity to change request timeout in current context + + :param timeout: + :return: + """ + timeout = self._prepare_timeout(timeout) + token = self._ctx_timeout.set(timeout) + try: + yield + finally: + self._ctx_timeout.reset(token) + async def close(self): """ Close all client sessions @@ -113,11 +160,11 @@ class BaseBot: :raise: :obj:`aiogram.exceptions.TelegramApiError` """ return await api.make_request(self.session, self.__token, method, data, files, - proxy=self.proxy, proxy_auth=self.proxy_auth, **kwargs) + proxy=self.proxy, proxy_auth=self.proxy_auth, timeout=self.timeout, **kwargs) async def download_file(self, file_path: base.String, destination: Optional[base.InputFile] = None, - timeout: Optional[base.Integer] = 30, + timeout: Optional[base.Integer] = sentinel, chunk_size: Optional[base.Integer] = 65536, seek: Optional[base.Boolean] = True) -> Union[io.BytesIO, io.FileIO]: """ diff --git a/aiogram/dispatcher/dispatcher.py b/aiogram/dispatcher/dispatcher.py index d1096718..0223c85f 100644 --- a/aiogram/dispatcher/dispatcher.py +++ b/aiogram/dispatcher/dispatcher.py @@ -5,6 +5,9 @@ import logging import time import typing +import aiohttp +from aiohttp.helpers import sentinel + from .filters import Command, ContentTypeFilter, ExceptionsFilter, FiltersFactory, FuncFilter, HashTag, Regexp, \ RegexpCommandsFilter, StateFilter, Text from .handler import Handler @@ -209,8 +212,13 @@ class Dispatcher(DataMixin, ContextInstanceMixin): return await self.bot.delete_webhook() - async def start_polling(self, timeout=20, relax=0.1, limit=None, reset_webhook=None, - fast: typing.Optional[bool] = True): + async def start_polling(self, + timeout=20, + relax=0.1, + limit=None, + reset_webhook=None, + fast: typing.Optional[bool] = True, + error_sleep: int = 5): """ Start long-polling @@ -238,12 +246,19 @@ class Dispatcher(DataMixin, ContextInstanceMixin): self._polling = True offset = None try: + current_request_timeout = self.bot.timeout + if current_request_timeout is not sentinel and timeout is not None: + request_timeout = aiohttp.ClientTimeout(total=current_request_timeout.total + timeout or 1) + else: + request_timeout = None + while self._polling: try: - updates = await self.bot.get_updates(limit=limit, offset=offset, timeout=timeout) + with self.bot.request_timeout(request_timeout): + updates = await self.bot.get_updates(limit=limit, offset=offset, timeout=timeout) except: log.exception('Cause exception while getting updates.') - await asyncio.sleep(15) + await asyncio.sleep(error_sleep) continue if updates: @@ -254,6 +269,7 @@ class Dispatcher(DataMixin, ContextInstanceMixin): if relax: await asyncio.sleep(relax) + finally: self._close_waiter._set_result(None) log.warning('Polling is stopped.') diff --git a/aiogram/utils/executor.py b/aiogram/utils/executor.py index 34acf6e9..65594371 100644 --- a/aiogram/utils/executor.py +++ b/aiogram/utils/executor.py @@ -23,7 +23,7 @@ def _setup_callbacks(executor, on_startup=None, on_shutdown=None): def start_polling(dispatcher, *, loop=None, skip_updates=False, reset_webhook=True, - on_startup=None, on_shutdown=None, timeout=None, fast=True): + on_startup=None, on_shutdown=None, timeout=20, fast=True): """ Start bot in long-polling mode @@ -291,7 +291,7 @@ class Executor: self.set_webhook(webhook_path=webhook_path, request_handler=request_handler, route_name=route_name) self.run_app(**kwargs) - def start_polling(self, reset_webhook=None, timeout=None, fast=True): + def start_polling(self, reset_webhook=None, timeout=20, fast=True): """ Start bot in long-polling mode