Add context middleware (with example of usage)

This commit is contained in:
Alex Root Junior 2018-01-10 20:39:56 +02:00
parent 9de31422eb
commit 2ab33fa1f8
3 changed files with 163 additions and 0 deletions

View file

@ -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()

View file

@ -35,6 +35,7 @@ class MiddlewareManager:
self.applications.append(middleware) self.applications.append(middleware)
middleware.setup(self) middleware.setup(self)
log.debug(f"Loaded middleware '{middleware.__class__.__name__}'") log.debug(f"Loaded middleware '{middleware.__class__.__name__}'")
return middleware
async def trigger(self, action: str, args: typing.Iterable): async def trigger(self, action: str, args: typing.Iterable):
""" """

View file

@ -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)