diff --git a/aiogram/dispatcher/handler.py b/aiogram/dispatcher/handler.py index 4ded9316..7be6b3d4 100644 --- a/aiogram/dispatcher/handler.py +++ b/aiogram/dispatcher/handler.py @@ -2,13 +2,14 @@ import inspect from contextvars import ContextVar ctx_data = ContextVar('ctx_handler_data') +current_handler = ContextVar('current_handler') -class SkipHandler(BaseException): +class SkipHandler(Exception): pass -class CancelHandler(BaseException): +class CancelHandler(Exception): pass @@ -87,9 +88,9 @@ class Handler: except FilterNotPassed: continue else: + ctx_token = current_handler.set(handler) try: if self.middleware_key: - # context.set_value('handler', handler) await self.dispatcher.middleware.trigger(f"process_{self.middleware_key}", args + (data,)) partial_data = _check_spec(handler, data) response = await handler(*args, **partial_data) @@ -101,6 +102,8 @@ class Handler: continue except CancelHandler: break + finally: + current_handler.reset(ctx_token) finally: if self.middleware_key: await self.dispatcher.middleware.trigger(f"post_process_{self.middleware_key}", diff --git a/examples/middleware_and_antiflood.py b/examples/middleware_and_antiflood.py index 7b83d9a4..6f26b2ee 100644 --- a/examples/middleware_and_antiflood.py +++ b/examples/middleware_and_antiflood.py @@ -3,8 +3,9 @@ import asyncio from aiogram import Bot, Dispatcher, executor, types from aiogram.contrib.fsm_storage.redis import RedisStorage2 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.utils.exceptions import Throttled TOKEN = 'BOT TOKEN HERE' @@ -52,11 +53,10 @@ class ThrottlingMiddleware(BaseMiddleware): :param message: """ # Get current handler - # handler = context.get_value('handler') + handler = current_handler.get() # Get dispatcher from context dispatcher = Dispatcher.current() - # If handler was configured, get rate limit and key from handler if handler: limit = getattr(handler, 'throttling_rate_limit', self.rate_limit) @@ -82,7 +82,7 @@ class ThrottlingMiddleware(BaseMiddleware): :param message: :param throttled: """ - # handler = context.get_value('handler') + handler = current_handler.get() dispatcher = Dispatcher.current() if handler: key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")