Skip to content

Commit aa0c29e

Browse files
lrdcxdesalissonlauffer
lrdcxdes
authored and
alissonlauffer
committed
feat: error handlers
1 parent ed4e775 commit aa0c29e

File tree

7 files changed

+202
-9
lines changed

7 files changed

+202
-9
lines changed

‎hydrogram/dispatcher.py‎

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
ChosenInlineResultHandler,
3232
DeletedMessagesHandler,
3333
EditedMessageHandler,
34+
ErrorHandler,
3435
InlineQueryHandler,
3536
MessageHandler,
3637
PollHandler,
@@ -79,6 +80,7 @@ def __init__(self, client: "hydrogram.Client"):
7980
self.updates_queue = asyncio.Queue()
8081
self.groups = OrderedDict()
8182
self._init_update_parsers()
83+
self.error_handlers = []
8284

8385
def _init_update_parsers(self):
8486
self.update_parsers = {
@@ -160,24 +162,38 @@ async def stop(self):
160162
await asyncio.gather(*self.handler_worker_tasks)
161163
self.handler_worker_tasks.clear()
162164
self.groups.clear()
165+
self.error_handlers.clear()
166+
163167
log.info("Stopped %s HandlerTasks", self.client.workers)
164168

165169
def add_handler(self, handler, group: int):
166170
async def fn():
167171
async with asyncio.Lock():
168-
if group not in self.groups:
169-
self.groups[group] = []
170-
self.groups = OrderedDict(sorted(self.groups.items()))
171-
self.groups[group].append(handler)
172+
if isinstance(handler, ErrorHandler):
173+
if handler not in self.error_handlers:
174+
self.error_handlers.append(handler)
175+
else:
176+
if group not in self.groups:
177+
self.groups[group] = []
178+
self.groups = OrderedDict(sorted(self.groups.items()))
179+
self.groups[group].append(handler)
172180

173181
self.loop.create_task(fn())
174182

175183
def remove_handler(self, handler, group: int):
176184
async def fn():
177185
async with asyncio.Lock():
178-
if group not in self.groups:
179-
raise ValueError(f"Group {group} does not exist. Handler was not removed.")
180-
self.groups[group].remove(handler)
186+
if isinstance(handler, ErrorHandler):
187+
if handler not in self.error_handlers:
188+
raise ValueError(
189+
f"Error handler {handler} does not exist. Handler was not removed."
190+
)
191+
192+
self.error_handlers.remove(handler)
193+
else:
194+
if group not in self.groups:
195+
raise ValueError(f"Group {group} does not exist. Handler was not removed.")
196+
self.groups[group].remove(handler)
181197

182198
self.loop.create_task(fn())
183199

@@ -225,7 +241,23 @@ async def _handle_update(self, handler, handler_type, parsed_update, update, use
225241
except hydrogram.ContinuePropagation:
226242
pass
227243
except Exception as e:
228-
log.exception(e)
244+
handled_error = False
245+
for error_handler in self.error_handlers:
246+
try:
247+
if await error_handler.check(self.client, e):
248+
await error_handler.callback(self.client, e)
249+
handled_error = True
250+
break
251+
except hydrogram.StopPropagation:
252+
raise
253+
except hydrogram.ContinuePropagation:
254+
continue
255+
except Exception as e:
256+
log.exception(e)
257+
continue
258+
259+
if not handled_error:
260+
log.exception(e)
229261

230262
async def _execute_callback(self, handler, *args):
231263
if inspect.iscoroutinefunction(handler.callback):

‎hydrogram/handlers/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from .deleted_messages_handler import DeletedMessagesHandler
2525
from .disconnect_handler import DisconnectHandler
2626
from .edited_message_handler import EditedMessageHandler
27+
from .error_handler import ErrorHandler
2728
from .inline_query_handler import InlineQueryHandler
2829
from .message_handler import MessageHandler
2930
from .poll_handler import PollHandler
@@ -38,6 +39,7 @@
3839
"DeletedMessagesHandler",
3940
"DisconnectHandler",
4041
"EditedMessageHandler",
42+
"ErrorHandler",
4143
"InlineQueryHandler",
4244
"MessageHandler",
4345
"PollHandler",

‎hydrogram/handlers/error_handler.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
# Hydrogram - Telegram MTProto API Client Library for Python
2+
# Copyright (C) 2023-present Hydrogram <https://hydrogram.org>
3+
#
4+
# This file is part of Hydrogram.
5+
#
6+
# Hydrogram is free software: you can redistribute it and/or modify
7+
# it under the terms of the GNU Lesser General Public License as published
8+
# by the Free Software Foundation, either version 3 of the License, or
9+
# (at your option) any later version.
10+
#
11+
# Hydrogram is distributed in the hope that it will be useful,
12+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
13+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14+
# GNU Lesser General Public License for more details.
15+
#
16+
# You should have received a copy of the GNU Lesser General Public License
17+
# along with Hydrogram. If not, see <http://www.gnu.org/licenses/>.
18+
19+
from typing import TYPE_CHECKING, Callable
20+
21+
from .handler import Handler
22+
23+
if TYPE_CHECKING:
24+
import hydrogram
25+
26+
27+
class ErrorHandler(Handler):
28+
"""The Error handler class. Used to handle errors.
29+
It is intended to be used with :meth:`~hydrogram.Client.add_handler`
30+
31+
For a nicer way to register this handler, have a look at the
32+
:meth:`~hydrogram.Client.on_message` decorator.
33+
34+
Parameters:
35+
callback (``Callable``):
36+
Pass a function that will be called when a new Error arrives. It takes *(client, error)*
37+
as positional arguments (look at the section below for a detailed description).
38+
39+
errors (:obj:`Exception` | List of :obj:`Exception`):
40+
Pass one or more exception classes to allow only a subset of errors to be passed
41+
in your callback function.
42+
43+
Other parameters:
44+
client (:obj:`~hydrogram.Client`):
45+
The Client itself, useful when you want to call other API methods inside the message handler.
46+
47+
error (:obj:`~Exception`):
48+
The error that was raised.
49+
50+
update (:obj:`~hydrogram.Update`):
51+
The update that caused the error.
52+
"""
53+
54+
def __init__(self, callback: Callable, errors=None):
55+
if errors is None:
56+
errors = [Exception]
57+
elif not isinstance(errors, list):
58+
errors = [errors]
59+
60+
self.errors = errors
61+
super().__init__(callback)
62+
63+
async def check(self, client: "hydrogram.Client", error: Exception):
64+
return any(isinstance(error, e) for e in self.errors)
65+
66+
def check_remove(self, error: Exception):
67+
return self.errors == error or any(isinstance(error, e) for e in self.errors)

‎hydrogram/methods/decorators/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,15 @@
2424
from .on_deleted_messages import OnDeletedMessages
2525
from .on_disconnect import OnDisconnect
2626
from .on_edited_message import OnEditedMessage
27+
from .on_error import OnError
2728
from .on_inline_query import OnInlineQuery
2829
from .on_message import OnMessage
2930
from .on_poll import OnPoll
3031
from .on_raw_update import OnRawUpdate
3132
from .on_user_status import OnUserStatus
3233

3334

34-
class Decorators(
35+
class Decorators( # noqa: N818 false-positive
3536
OnMessage,
3637
OnEditedMessage,
3738
OnDeletedMessages,
@@ -44,5 +45,6 @@ class Decorators(
4445
OnChosenInlineResult,
4546
OnChatMemberUpdated,
4647
OnChatJoinRequest,
48+
OnError,
4749
):
4850
pass

‎hydrogram/methods/decorators/on_error.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Hydrogram - Telegram MTProto API Client Library for Python
2+
# Copyright (C) 2023-present Hydrogram <https://hydrogram.org>
3+
#
4+
# This file is part of Hydrogram.
5+
#
6+
# Hydrogram is free software: you can redistribute it and/or modify
7+
# it under the terms of the GNU Lesser General Public License as published
8+
# by the Free Software Foundation, either version 3 of the License, or
9+
# (at your option) any later version.
10+
#
11+
# Hydrogram is distributed in the hope that it will be useful,
12+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
13+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14+
# GNU Lesser General Public License for more details.
15+
#
16+
# You should have received a copy of the GNU Lesser General Public License
17+
# along with Hydrogram. If not, see <http://www.gnu.org/licenses/>.
18+
19+
from typing import Callable
20+
21+
import hydrogram
22+
from hydrogram.filters import Filter
23+
24+
25+
class OnError:
26+
def on_error(self=None, errors=None) -> Callable:
27+
"""Decorator for handling new errors.
28+
29+
This does the same thing as :meth:`~hydrogram.Client.add_handler` using the
30+
:obj:`~hydrogram.handlers.MessageHandler`.
31+
32+
Parameters:
33+
errors (:obj:`~Exception`, *optional*):
34+
Pass one or more errors to allow only a subset of errors to be passed
35+
in your function.
36+
"""
37+
38+
def decorator(func: Callable) -> Callable:
39+
if isinstance(self, hydrogram.Client):
40+
self.add_handler(hydrogram.handlers.ErrorHandler(func, errors), 0)
41+
elif isinstance(self, Filter) or self is None:
42+
if not hasattr(func, "handlers"):
43+
func.handlers = []
44+
45+
func.handlers.append((hydrogram.handlers.ErrorHandler(func, self), 0))
46+
47+
return func
48+
49+
return decorator

‎hydrogram/methods/utilities/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
from .add_handler import AddHandler
2121
from .export_session_string import ExportSessionString
22+
from .remove_error_handler import RemoveErrorHandler
2223
from .remove_handler import RemoveHandler
2324
from .restart import Restart
2425
from .run import Run
@@ -31,6 +32,7 @@ class Utilities(
3132
AddHandler,
3233
ExportSessionString,
3334
RemoveHandler,
35+
RemoveErrorHandler,
3436
Restart,
3537
Run,
3638
Start,

‎hydrogram/methods/utilities/remove_error_handler.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# Hydrogram - Telegram MTProto API Client Library for Python
2+
# Copyright (C) 2023-present Hydrogram <https://hydrogram.org>
3+
#
4+
# This file is part of Hydrogram.
5+
#
6+
# Hydrogram is free software: you can redistribute it and/or modify
7+
# it under the terms of the GNU Lesser General Public License as published
8+
# by the Free Software Foundation, either version 3 of the License, or
9+
# (at your option) any later version.
10+
#
11+
# Hydrogram is distributed in the hope that it will be useful,
12+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
13+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14+
# GNU Lesser General Public License for more details.
15+
#
16+
# You should have received a copy of the GNU Lesser General Public License
17+
# along with Hydrogram. If not, see <http://www.gnu.org/licenses/>.
18+
19+
from __future__ import annotations
20+
21+
from typing import TYPE_CHECKING
22+
23+
if TYPE_CHECKING:
24+
import hydrogram
25+
26+
27+
class RemoveErrorHandler:
28+
def remove_error_handler(
29+
self: hydrogram.Client, error: Exception | tuple[Exception] = Exception
30+
):
31+
"""Remove a previously-registered error handler. (using exception classes)
32+
33+
Parameters:
34+
error (``Exception``):
35+
The error(s) for handlers to be removed.
36+
"""
37+
for handler in self.dispatcher.error_handlers:
38+
if handler.check_remove(error):
39+
self.dispatcher.error_handlers.remove(handler)

0 commit comments

Comments
 (0)