From 03c7e4a6d868b81743588957d92708ed30e2b0d5 Mon Sep 17 00:00:00 2001
From: Alex Root Junior <jroot.junior@gmail.com>
Date: Sat, 3 Jun 2017 10:50:48 +0300
Subject: [PATCH] Provide to cancel next message waiter.

---
 aiogram/dispatcher/__init__.py         |  4 ++--
 aiogram/dispatcher/filters.py          | 11 +++++++++++
 aiogram/dispatcher/handler.py          | 26 ++++++++++++++++++++------
 examples/steps_and_downloading_file.py |  7 ++++++-
 4 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/aiogram/dispatcher/__init__.py b/aiogram/dispatcher/__init__.py
index 4176df4d..6156c790 100644
--- a/aiogram/dispatcher/__init__.py
+++ b/aiogram/dispatcher/__init__.py
@@ -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)
diff --git a/aiogram/dispatcher/filters.py b/aiogram/dispatcher/filters.py
index c8eebb8e..4907a802 100644
--- a/aiogram/dispatcher/filters.py
+++ b/aiogram/dispatcher/filters.py
@@ -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 = []
 
diff --git a/aiogram/dispatcher/handler.py b/aiogram/dispatcher/handler.py
index 7afa3659..a1ff7b12 100644
--- a/aiogram/dispatcher/handler.py
+++ b/aiogram/dispatcher/handler.py
@@ -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
diff --git a/examples/steps_and_downloading_file.py b/examples/steps_and_downloading_file.py
index 9e2f720b..a2e229e9 100644
--- a/examples/steps_and_downloading_file.py
+++ b/examples/steps_and_downloading_file.py
@@ -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)