python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2026-04-08 16:37:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Added timeout for every method.

Browse source
This commit is contained in:
Forden 2019-03-21 02:04:51 +03:00
parent 39c8d859dc
commit a8366125f1
6 changed files with 375 additions and 199 deletions
Download patch file Download diff file Expand all files Collapse all files

21
aiogram/bot/api.py
View file

@ -52,7 +52,8 @@ async def check_result(method_name: str, content_type: str, status_code: int, bo
log.debug('Response for %s: [%d] "%r"', method_name, status_code, body)
if content_type != 'application/json':
raise exceptions.NetworkError(f"Invalid response with content type {content_type}: \"{body}\"")
raise exceptions.NetworkError(
f"Invalid response with content type {content_type}: \"{body}\"")
try:
result_json = json.loads(body)
@ -60,7 +61,8 @@ async def check_result(method_name: str, content_type: str, status_code: int, bo
result_json = {}
description = result_json.get('description') or body
parameters = types.ResponseParameters(**result_json.get('parameters', {}) or {})
parameters = types.ResponseParameters(
**result_json.get('parameters', {}) or {})
if HTTPStatus.OK <= status_code <= HTTPStatus.IM_USED:
return result_json.get('result')
@ -86,18 +88,22 @@ async def check_result(method_name: str, content_type: str, status_code: int, bo
raise exceptions.TelegramAPIError(f"{description} [{status_code}]")
async def make_request(session, token, method, data=None, files=None, **kwargs):
async def make_request(session, token, method, data=None, files=None, timeout=None, **kwargs):
# log.debug(f"Make request: '{method}' with data: {data} and files {files}")
log.debug('Make request: "%s" with data: "%r" and files "%r"', method, data, files)
log.debug('Make request: "%s" with data: "%r" and files "%r"',
method, data, files)
url = Methods.api_url(token=token, method=method)
req = compose_data(data, files)
timeout = aiohttp.ClientTimeout(timeout)
try:
async with session.post(url, data=req, **kwargs) as response:
async with session.post(url, data=req, timeout=timeout, **kwargs) as response:
return await check_result(method, response.content_type, response.status, await response.text())
except aiohttp.ClientError as e:
raise exceptions.NetworkError(f"aiohttp client throws an error: {e.__class__.__name__}: {e}")
raise exceptions.NetworkError(
f"aiohttp client throws an error: {e.__class__.__name__}: {e}")
def guess_filename(obj):
@ -132,7 +138,8 @@ def compose_data(params=None, files=None):
if len(f) == 2:
filename, fileobj = f
else:
raise ValueError('Tuple must have exactly 2 elements: filename, fileobj')
raise ValueError(
'Tuple must have exactly 2 elements: filename, fileobj')
elif isinstance(f, types.InputFile):
filename, fileobj = f.filename, f.file
else:

18
aiogram/bot/base.py
View file

@ -20,8 +20,10 @@ class BaseBot:
def __init__(self, token: base.String,
loop: Optional[Union[asyncio.BaseEventLoop, asyncio.AbstractEventLoop]] = None,
connections_limit: Optional[base.Integer] = None,
proxy: Optional[base.String] = None, proxy_auth: Optional[aiohttp.BasicAuth] = None,
proxy: Optional[base.String] = None,
proxy_auth: Optional[aiohttp.BasicAuth] = None,
validate_token: Optional[base.Boolean] = True,
connection_timeout: Optional[base.Integer] = None,
parse_mode=None):
"""
Instructions how to get Bot token is found here: https://core.telegram.org/bots#3-how-do-i-create-a-bot
@ -49,6 +51,7 @@ class BaseBot:
self.proxy = proxy
self.proxy_auth = proxy_auth
self.connection_timeout = aiohttp.ClientTimeout(total=connection_timeout)
# Asyncio loop instance
if loop is None:
@ -79,7 +82,7 @@ class BaseBot:
else:
connector = aiohttp.TCPConnector(limit=connections_limit, ssl=ssl_context, loop=self.loop)
self.session = aiohttp.ClientSession(connector=connector, loop=self.loop, json_serialize=json.dumps)
self.session = aiohttp.ClientSession(connector=connector, loop=self.loop, json_serialize=json.dumps, timeout=self.connection_timeout)
self.parse_mode = parse_mode
@ -91,7 +94,8 @@ class BaseBot:
async def request(self, method: base.String,
data: Optional[Dict] = None,
files: Optional[Dict] = None, **kwargs) -> Union[List, Dict, base.Boolean]:
files: Optional[Dict] = None,
timeout: Optional[int] = None, **kwargs) -> Union[List, Dict, base.Boolean]:
"""
Make an request to Telegram Bot API
@ -108,7 +112,7 @@ 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=timeout, **kwargs)
async def download_file(self, file_path: base.String,
destination: Optional[base.InputFile] = None,
@ -134,7 +138,8 @@ class BaseBot:
url = api.Methods.file_url(token=self.__token, path=file_path)
dest = destination if isinstance(destination, io.IOBase) else open(destination, 'wb')
dest = destination if isinstance(
destination, io.IOBase) else open(destination, 'wb')
async with self.session.get(url, timeout=timeout, proxy=self.proxy, proxy_auth=self.proxy_auth) as response:
while True:
chunk = await response.content.read(chunk_size)
@ -182,7 +187,8 @@ class BaseBot:
raise TypeError(f"Parse mode must be str, not {type(value)}")
value = value.lower()
if value not in ParseMode.all():
raise ValueError(f"Parse mode must be one of {ParseMode.all()}")
raise ValueError(
f"Parse mode must be one of {ParseMode.all()}")
setattr(self, '_parse_mode', value)
@parse_mode.deleter

428
aiogram/bot/bot.py
View file

File diff suppressed because it is too large Load diff

39
aiogram/types/message.py
View file

@ -51,7 +51,8 @@ class Message(base.TelegramObject):
author_signature: base.String = fields.Field()
text: base.String = fields.Field()
entities: typing.List[MessageEntity] = fields.ListField(base=MessageEntity)
caption_entities: typing.List[MessageEntity] = fields.ListField(base=MessageEntity)
caption_entities: typing.List[MessageEntity] = fields.ListField(
base=MessageEntity)
audio: Audio = fields.Field(base=Audio)
document: Document = fields.Field(base=Document)
animation: Animation = fields.Field(base=Animation)
@ -77,7 +78,8 @@ class Message(base.TelegramObject):
migrate_from_chat_id: base.Integer = fields.Field()
pinned_message: Message = fields.Field(base='Message')
invoice: Invoice = fields.Field(base=Invoice)
successful_payment: SuccessfulPayment = fields.Field(base=SuccessfulPayment)
successful_payment: SuccessfulPayment = fields.Field(
base=SuccessfulPayment)
connected_website: base.String = fields.Field()
passport_data: PassportData = fields.Field(base=PassportData)
@ -192,7 +194,7 @@ class Message(base.TelegramObject):
raise TypeError("This message doesn't have any text.")
quote_fn = md.quote_html if as_html else md.escape_md
entities = self.entities or self.caption_entities
if not entities:
return quote_fn(text)
@ -293,7 +295,8 @@ class Message(base.TelegramObject):
disable_web_page_preview=disable_web_page_preview,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def reply_photo(self, photo: typing.Union[base.InputFile, base.String],
caption: typing.Union[base.String, None] = None,
@ -320,7 +323,8 @@ class Message(base.TelegramObject):
return await self.bot.send_photo(chat_id=self.chat.id, photo=photo, caption=caption,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def reply_audio(self, audio: typing.Union[base.InputFile, base.String],
caption: typing.Union[base.String, None] = None,
@ -365,7 +369,8 @@ class Message(base.TelegramObject):
title=title,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def send_animation(self,
animation: typing.Union[base.InputFile, base.String],
@ -423,8 +428,8 @@ class Message(base.TelegramObject):
parse_mode=parse_mode,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup
)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def reply_document(self, document: typing.Union[base.InputFile, base.String],
caption: typing.Union[base.String, None] = None,
@ -456,7 +461,8 @@ class Message(base.TelegramObject):
caption=caption,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def reply_video(self, video: typing.Union[base.InputFile, base.String],
duration: typing.Union[base.Integer, None] = None,
@ -499,7 +505,8 @@ class Message(base.TelegramObject):
caption=caption,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def reply_voice(self, voice: typing.Union[base.InputFile, base.String],
caption: typing.Union[base.String, None] = None,
@ -537,7 +544,8 @@ class Message(base.TelegramObject):
duration=duration,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def reply_video_note(self, video_note: typing.Union[base.InputFile, base.String],
duration: typing.Union[base.Integer, None] = None,
@ -572,7 +580,8 @@ class Message(base.TelegramObject):
length=length,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def reply_media_group(self, media: typing.Union[MediaGroup, typing.List],
disable_notification: typing.Union[base.Boolean, None] = None,
@ -593,7 +602,8 @@ class Message(base.TelegramObject):
return await self.bot.send_media_group(self.chat.id,
media=media,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None)
reply_to_message_id=self.message_id if reply else None,
timeout=self.bot.connection_timeout.total)
async def reply_location(self, latitude: base.Float,
longitude: base.Float, live_period: typing.Union[base.Integer, None] = None,
@ -626,7 +636,8 @@ class Message(base.TelegramObject):
live_period=live_period,
disable_notification=disable_notification,
reply_to_message_id=self.message_id if reply else None,
reply_markup=reply_markup)
reply_markup=reply_markup,
timeout=self.bot.connection_timeout.total)
async def edit_live_location(self, latitude: base.Float, longitude: base.Float,
reply_markup=None) -> typing.Union[Message, base.Boolean]:

21
examples/timeout_example.py Normal file
View file

@ -0,0 +1,21 @@
import asyncio
import logging
from aiogram import Bot, types, filters
from aiogram.dispatcher import Dispatcher
from aiogram.utils.executor import start_polling
API_TOKEN = 'BOT TOKEN HERE'
logging.basicConfig(level=logging.INFO)
loop = asyncio.get_event_loop()
bot = Bot(token=API_TOKEN, loop=loop, connection_timeout=5)
dp = Dispatcher(bot)
@dp.message_handler(filters.CommandStart())
async def start(message: types.Message):
await message.reply_photo(types.InputFile('data/cat.jpg'), f'Cat with Bot\'s timeout: {bot.connection_timeout.total}!')
await bot.send_photo(message.chat.id, types.InputFile('data/cats.jpg'), 'More cats with timeout 1 second!', timeout=1)
if __name__ == '__main__':
start_polling(dp, loop=loop, skip_updates=True)

47
tests/test_message.py
View file

@ -1,47 +0,0 @@
from asyncio import BaseEventLoop
import pytest
from aiogram import Bot, types
from . import FakeTelegram, TOKEN
pytestmark = pytest.mark.asyncio
@pytest.yield_fixture()
async def bot(event_loop):
""" Bot fixture """
_bot = Bot(TOKEN, loop=event_loop, parse_mode=types.ParseMode.HTML)
yield _bot
await _bot.close()
@pytest.yield_fixture()
async def message(bot, event_loop):
"""
Message fixture
:param bot: Telegram bot fixture
:type bot: Bot
:param event_loop: asyncio event loop
:type event_loop: BaseEventLoop
"""
from .types.dataset import MESSAGE
msg = types.Message(**MESSAGE)
async with FakeTelegram(message_dict=MESSAGE, loop=event_loop):
_message = await bot.send_message(chat_id=msg.chat.id, text=msg.text)
yield _message
class TestMiscCases:
async def test_calling_bot_not_from_context(self, message):
"""
Calling any helper method without bot instance in context.
:param message: message fixture
:type message: types.Message
:return: RuntimeError with reason and help
"""
with pytest.raises(RuntimeError):
await message.edit_text('test_calling_bot_not_from_context')