Provide to cancel next message waiter.

This commit is contained in:
Alex Root Junior 2017-06-03 10:50:48 +03:00
parent bf9657dfc5
commit 03c7e4a6d8
4 changed files with 39 additions and 9 deletions

View file

@ -255,7 +255,7 @@ class Dispatcher:
return decorator
async def next_message(self, message: types.Message, otherwise=None, once=False,
async def next_message(self, message: types.Message, otherwise=None, once=False, include_cancel=True,
regexp=None, content_types=None, func=None, custom_filters=None, **kwargs):
if content_types is None:
content_types = []
@ -267,5 +267,5 @@ class Dispatcher:
content_types=content_types,
func=func,
**kwargs)
self.next_step_message_handlers.register(message, otherwise, once, filters_set)
self.next_step_message_handlers.register(message, otherwise, once, include_cancel, filters_set)
return await self.next_step_message_handlers.wait(message)

View file

@ -80,6 +80,17 @@ class ContentTypeFilter(Filter):
return message.content_type in self.content_types
class CancelFilter(Filter):
def __init__(self, cancel_set=None):
if cancel_set is None:
cancel_set = ['/cancel', 'cancel', 'cancel.']
self.cancel_set = cancel_set
def check(self, message):
if message.text:
return message.text.lower() in self.cancel_set
def generate_default_filters(*args, **kwargs):
filters_set = []

View file

@ -1,7 +1,7 @@
from asyncio import Event
from aiogram import types
from .filters import check_filters
from .filters import check_filters, CancelFilter
from .. import types
class Handler:
@ -37,13 +37,17 @@ class NextStepHandler:
self.dispatcher = dispatcher
self.handlers = {}
def register(self, message, otherwise=None, once=False, filters=None):
def register(self, message, otherwise=None, once=False, include_cancel=False, filters=None):
identifier = gen_identifier(message.chat.id, message.chat.id)
if identifier not in self.handlers:
self.handlers[identifier] = {'event': Event(), 'filters': filters,
'otherwise': otherwise, 'once': once}
'otherwise': otherwise, 'once': once,
'include_cancel': include_cancel,
'message': None}
return True
# In normal it's impossible.
raise RuntimeError('Dialog already wait message.')
# return False
@ -52,12 +56,22 @@ class NextStepHandler:
if identifier not in self.handlers:
return False
handler = self.handlers[identifier]
include_cancel = handler['include_cancel']
if include_cancel:
filter_ = CancelFilter(include_cancel if isinstance(include_cancel, (list, set, tuple)) else None)
if filter_.check(message):
handler['event'].set()
return True
if handler['filters'] and not await check_filters(handler['filters'], [message], {}):
otherwise = handler['otherwise']
if otherwise:
await otherwise(message)
if not handler['once']:
return False
if handler['once']:
handler['event'].set()
return True
handler['message'] = message
handler['event'].set()
return True

View file

@ -37,7 +37,12 @@ async def save_sticker(message):
# It can only be a sticker
msg = await dp.next_message(message,
content_types=[ContentType.STICKER],
otherwise=handle_bad_message)
otherwise=handle_bad_message,
include_cancel=True)
if not msg:
# If user send /cancel
return await message.reply('Canceled.')
# Download file to memory (io.BytesIO)
photo = await bot.download_file(msg.sticker.file_id)