From 2ab33fa1f8a5884e21a84fe00c73e7da3acaafa9 Mon Sep 17 00:00:00 2001
From: Alex Root Junior <jroot.junior@gmail.com>
Date: Wed, 10 Jan 2018 20:39:56 +0200
Subject: [PATCH] Add context middleware (with example of usage)

---
 aiogram/contrib/middlewares/context.py | 115 +++++++++++++++++++++++++
 aiogram/dispatcher/middlewares.py      |   1 +
 examples/example_context_middleware.py |  47 ++++++++++
 3 files changed, 163 insertions(+)
 create mode 100644 aiogram/contrib/middlewares/context.py
 create mode 100644 examples/example_context_middleware.py

diff --git a/aiogram/contrib/middlewares/context.py b/aiogram/contrib/middlewares/context.py
new file mode 100644
index 00000000..8e6dce7a
--- /dev/null
+++ b/aiogram/contrib/middlewares/context.py
@@ -0,0 +1,115 @@
+from aiogram import types
+from aiogram.dispatcher import ctx
+from aiogram.dispatcher.middlewares import BaseMiddleware
+
+OBJ_KEY = '_context_data'
+
+
+class ContextMiddleware(BaseMiddleware):
+    """
+    Allow to store data at all of lifetime of Update object
+    """
+
+    async def on_pre_process_update(self, update: types.Update):
+        """
+        Start of Update lifetime
+
+        :param update:
+        :return:
+        """
+        self._configure_update(update)
+
+    async def on_post_process_update(self, update: types.Update, result):
+        """
+        On finishing of processing update
+
+        :param update:
+        :param result:
+        :return:
+        """
+        if OBJ_KEY in update.conf:
+            del update.conf[OBJ_KEY]
+
+    def _configure_update(self, update: types.Update = None):
+        """
+        Setup data storage
+
+        :param update:
+        :return:
+        """
+        obj = update.conf[OBJ_KEY] = {}
+        return obj
+
+    def _get_dict(self):
+        """
+        Get data from update stored in current context
+
+        :return:
+        """
+        update = ctx.get_update()
+        obj = update.conf.get(OBJ_KEY, None)
+        if obj is None:
+            obj = self._configure_update(update)
+        return obj
+
+    def __getitem__(self, item):
+        """
+        Item getter
+
+        :param item:
+        :return:
+        """
+        return self._get_dict()[item]
+
+    def __setitem__(self, key, value):
+        """
+        Item setter
+
+        :param key:
+        :param value:
+        :return:
+        """
+        data = self._get_dict()
+        data[key] = value
+
+    def __iter__(self):
+        """
+        Iterate over dict
+
+        :return:
+        """
+        return self._get_dict().__iter__()
+
+    def keys(self):
+        """
+        Iterate over dict keys
+
+        :return:
+        """
+        return self._get_dict().keys()
+
+    def values(self):
+        """
+        Iterate over dict values
+
+        :return:
+        """
+        return self._get_dict().values()
+
+    def get(self, key, default=None):
+        """
+        Get item from dit or return default value
+
+        :param key:
+        :param default:
+        :return:
+        """
+        return self._get_dict().get(key, default)
+
+    def export(self):
+        """
+        Export all data s dict
+
+        :return:
+        """
+        return self._get_dict()
diff --git a/aiogram/dispatcher/middlewares.py b/aiogram/dispatcher/middlewares.py
index 85c7ef2d..2d8f11a6 100644
--- a/aiogram/dispatcher/middlewares.py
+++ b/aiogram/dispatcher/middlewares.py
@@ -35,6 +35,7 @@ class MiddlewareManager:
         self.applications.append(middleware)
         middleware.setup(self)
         log.debug(f"Loaded middleware '{middleware.__class__.__name__}'")
+        return middleware
 
     async def trigger(self, action: str, args: typing.Iterable):
         """
diff --git a/examples/example_context_middleware.py b/examples/example_context_middleware.py
new file mode 100644
index 00000000..9f6be968
--- /dev/null
+++ b/examples/example_context_middleware.py
@@ -0,0 +1,47 @@
+from aiogram import Bot, types
+from aiogram.contrib.middlewares.context import ContextMiddleware
+from aiogram.dispatcher import Dispatcher
+from aiogram.types import ParseMode
+from aiogram.utils import markdown as md
+from aiogram.utils.executor import start_polling
+
+API_TOKEN = 'BOT TOKEN HERE'
+
+bot = Bot(token=API_TOKEN)
+dp = Dispatcher(bot)
+
+# Setup Context middleware
+data: ContextMiddleware = dp.middleware.setup(ContextMiddleware())
+
+
+# Write custom filter
+async def demo_filter(message: types.Message):
+    # Store some data in context
+    command = data['command'] = message.get_command() or ''
+    args = data['args'] = message.get_args() or ''
+    data['has_args'] = bool(args)
+    data['some_random_data'] = 42
+    return command != '/bad_command'
+
+
+@dp.message_handler(demo_filter)
+async def send_welcome(message: types.Message):
+    # Get data from context
+    # All of that available only in current context and from current update object
+    # `data`- pseudo-alias for `ctx.get_update().conf['_context_data']`
+    command = data['command']
+    args = data['args']
+    rand = data['some_random_data']
+    has_args = data['has_args']
+
+    # Send as pre-formatted code block.
+    await message.reply(md.hpre(f"""command: {command}
+args: {['Not available', 'available'][has_args]}: {args}
+some random data: {rand}
+message ID: {message.message_id}
+message: {message.html_text}
+    """), parse_mode=ParseMode.HTML)
+
+
+if __name__ == '__main__':
+    start_polling(dp)