python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2025-12-11 18:01:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Allow to get current handler

Browse source
This commit is contained in:
Alex Root Junior 2018-09-21 22:34:59 +03:00
parent 5aff7f7d35
commit 5be3f9d00b
2 changed files with 10 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

9
aiogram/dispatcher/handler.py
View file

@ -2,13 +2,14 @@ import inspect
from contextvars import ContextVar from contextvars import ContextVar
ctx_data = ContextVar('ctx_handler_data') ctx_data = ContextVar('ctx_handler_data')
current_handler = ContextVar('current_handler')
class SkipHandler(BaseException): class SkipHandler(Exception):
pass pass
class CancelHandler(BaseException): class CancelHandler(Exception):
pass pass
@ -87,9 +88,9 @@ class Handler:
except FilterNotPassed: except FilterNotPassed:
continue continue
else: else:
ctx_token = current_handler.set(handler)
try: try:
if self.middleware_key: if self.middleware_key:
# context.set_value('handler', handler)
await self.dispatcher.middleware.trigger(f"process_{self.middleware_key}", args + (data,)) await self.dispatcher.middleware.trigger(f"process_{self.middleware_key}", args + (data,))
partial_data = _check_spec(handler, data) partial_data = _check_spec(handler, data)
response = await handler(*args, **partial_data) response = await handler(*args, **partial_data)
@ -101,6 +102,8 @@ class Handler:
continue continue
except CancelHandler: except CancelHandler:
break break
finally:
current_handler.reset(ctx_token)
finally: finally:
if self.middleware_key: if self.middleware_key:
await self.dispatcher.middleware.trigger(f"post_process_{self.middleware_key}", await self.dispatcher.middleware.trigger(f"post_process_{self.middleware_key}",

8
examples/middleware_and_antiflood.py
View file

@ -3,8 +3,9 @@ import asyncio
from aiogram import Bot, Dispatcher, executor, types from aiogram import Bot, Dispatcher, executor, types
from aiogram.contrib.fsm_storage.redis import RedisStorage2 from aiogram.contrib.fsm_storage.redis import RedisStorage2
from aiogram.dispatcher import DEFAULT_RATE_LIMIT from aiogram.dispatcher import DEFAULT_RATE_LIMIT
from aiogram.dispatcher.handler import CancelHandler from aiogram.dispatcher.handler import CancelHandler, current_handler
from aiogram.dispatcher.middlewares import BaseMiddleware from aiogram.dispatcher.middlewares import BaseMiddleware
from aiogram.utils.exceptions import Throttled
TOKEN = 'BOT TOKEN HERE' TOKEN = 'BOT TOKEN HERE'
@ -52,11 +53,10 @@ class ThrottlingMiddleware(BaseMiddleware):
:param message: :param message:
""" """
# Get current handler # Get current handler
# handler = context.get_value('handler') handler = current_handler.get()
# Get dispatcher from context # Get dispatcher from context
dispatcher = Dispatcher.current() dispatcher = Dispatcher.current()
# If handler was configured, get rate limit and key from handler # If handler was configured, get rate limit and key from handler
if handler: if handler:
limit = getattr(handler, 'throttling_rate_limit', self.rate_limit) limit = getattr(handler, 'throttling_rate_limit', self.rate_limit)
@ -82,7 +82,7 @@ class ThrottlingMiddleware(BaseMiddleware):
:param message: :param message:
:param throttled: :param throttled:
""" """
# handler = context.get_value('handler') handler = current_handler.get()
dispatcher = Dispatcher.current() dispatcher = Dispatcher.current()
if handler: if handler:
key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}") key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")