From b7bba77d299d9f5a144f4ea1e551703176f197fb Mon Sep 17 00:00:00 2001 From: Alex Root Junior Date: Sat, 27 May 2017 03:09:14 +0300 Subject: [PATCH] Add base of messages handler. --- aiogram/dispatcher/__init__.py | 51 +++++++++++++++++++++- aiogram/dispatcher/filters.py | 80 ++++++++++++++++++++++++++++++++++ aiogram/dispatcher/handler.py | 25 +++-------- 3 files changed, 134 insertions(+), 22 deletions(-) create mode 100644 aiogram/dispatcher/filters.py diff --git a/aiogram/dispatcher/__init__.py b/aiogram/dispatcher/__init__.py index d99775a8..405565f4 100644 --- a/aiogram/dispatcher/__init__.py +++ b/aiogram/dispatcher/__init__.py @@ -1,15 +1,21 @@ import asyncio import logging +from .filters import CommandsFilter, RegexpFilter, ContentTypeFilter from .handler import Handler from ..bot import AIOGramBot +from ..types.message import ContentType log = logging.getLogger(__name__) class Dispatcher: - def __init__(self, bot): + def __init__(self, bot, loop=None): self.bot: AIOGramBot = bot + if loop is None: + loop = self.bot.loop + + self.loop = loop self.last_update_id = 0 self.updates = Handler(self) @@ -33,7 +39,7 @@ class Dispatcher: async def process_updates(self, updates): for update in updates: - self.bot.loop.create_task(self.updates.notify(update)) + self.loop.create_task(self.updates.notify(update)) async def process_update(self, update): if update.message: @@ -42,6 +48,8 @@ class Dispatcher: async def start_pooling(self, timeout=20, relax=0.1): if self._pooling: raise RuntimeError('Pooling already started') + log.info('Start pooling.') + self._pooling = True offset = None while self._pooling: @@ -53,10 +61,49 @@ class Dispatcher: continue if updates: + log.info(f"Received {len(updates)} updates.") offset = updates[-1].update_id + 1 await self.process_updates(updates) await asyncio.sleep(relax) + log.warning('Pooling is stopped.') + def stop_pooling(self): self._pooling = False + + def message_handler(self, commands=None, regexp=None, content_type=None, func=None, + custom_filters=None): + if commands is None: + commands = [] + if content_type is None: + content_type = [ContentType.TEXT] + if custom_filters is None: + custom_filters = [] + + filters_preset = [] + if commands: + if isinstance(commands, str): + commands = [commands] + filters_preset.append(CommandsFilter(commands)) + + if regexp: + filters_preset.append(RegexpFilter(regexp)) + + if content_type: + filters_preset.append(ContentTypeFilter(content_type)) + + if func: + filters_preset.append(func) + + if custom_filters: + filters_preset += custom_filters + + def decorator(func): + self.messages.register(func, filters_preset) + return func + + return decorator + + def __del__(self): + self._pooling = False diff --git a/aiogram/dispatcher/filters.py b/aiogram/dispatcher/filters.py new file mode 100644 index 00000000..2f3a1a6a --- /dev/null +++ b/aiogram/dispatcher/filters.py @@ -0,0 +1,80 @@ +import inspect +import re + + +async def check_filter(filter_, args, kwargs): + if any((inspect.isasyncgen(filter_), + inspect.iscoroutine(filter_), + inspect.isawaitable(filter_), + inspect.isasyncgenfunction(filter_), + inspect.iscoroutinefunction(filter_))): + return await filter_(*args, **kwargs) + elif callable(filter_): + return filter_(*args, **kwargs) + else: + return True + + +async def check_filters(filters, args, kwargs): + if filters is not None: + for filter_ in filters: + f = await check_filter(filter_, args, kwargs) + if not f: + return False + return True + + +class Filter: + def __call__(self, *args, **kwargs): + return self.check(*args, **kwargs) + + def check(self, *args, **kwargs): + raise NotImplementedError + + +class AsyncFilter(Filter): + def __aiter__(self): + return None + + def __await__(self): + return self.check + + async def check(self, *args, **kwargs): + pass + + +class CommandsFilter(AsyncFilter): + def __init__(self, commands): + self.commands = commands + + async def check(self, message): + if not message.is_command(): + return False + + command = message.text.split()[0][1:] + command, _, mention = command.partition('@') + + if mention and mention != (await message.bot.me).username: + return False + + if command not in self.commands: + return False + + return True + + +class RegexpFilter(Filter): + def __init__(self, regexp): + self.regexp = re.compile(regexp) + + def check(self, message): + if message.text: + return bool(self.regexp.match(message.text)) + + +class ContentTypeFilter(Filter): + def __init__(self, content_types): + self.content_types = content_types + + def check(self, message): + return message.content_type in self.content_types diff --git a/aiogram/dispatcher/handler.py b/aiogram/dispatcher/handler.py index 006c199e..2090fe1a 100644 --- a/aiogram/dispatcher/handler.py +++ b/aiogram/dispatcher/handler.py @@ -1,27 +1,10 @@ -import inspect - - -async def check_filter(filter_, args, kwargs): - if inspect.iscoroutinefunction(filter_): - return await filter_(*args, **kwargs) - elif callable(filter_): - return filter_(*args, **kwargs) - else: - return True - - -async def check_filters(filters, args, kwargs): - if filters is not None: - for filter_ in filters: - f = await check_filter(filter_, args, kwargs) - if not f: - return False - return True +from .filters import check_filters class Handler: - def __init__(self, dispatcher): + def __init__(self, dispatcher, once=True): self.dispatcher = dispatcher + self.once = once self.handlers = [] @@ -42,3 +25,5 @@ class Handler: for filters, handler in self.handlers: if await check_filters(filters, args, kwargs): await handler(*args, **kwargs) + if self.once: + break