Add PollHandler type and on_poll decorator for handling Poll updates

delivrance · delivrance · commit 5905f761fa21 · 2019-04-14T20:50:13.000+02:00
diff --git a/pyrogram/client/ext/dispatcher.py b/pyrogram/client/ext/dispatcher.py
@@ -26,7 +26,7 @@
 from pyrogram.api import types
 from ..handlers import (
     CallbackQueryHandler, MessageHandler, DeletedMessagesHandler,
-    UserStatusHandler, RawUpdateHandler, InlineQueryHandler
+    UserStatusHandler, RawUpdateHandler, InlineQueryHandler, PollHandler
 )
 
 log = logging.getLogger(__name__)
@@ -79,7 +79,10 @@ def __init__(self, client, workers: int):
                 ),
 
             (types.UpdateBotInlineQuery,):
-                lambda upd, usr, cht: (pyrogram.InlineQuery._parse(self.client, upd, usr), InlineQueryHandler)
+                lambda upd, usr, cht: (pyrogram.InlineQuery._parse(self.client, upd, usr), InlineQueryHandler),
+
+            (types.UpdateMessagePoll,):
+                lambda upd, usr, cht: (pyrogram.Poll._parse(self.client, upd), PollHandler)
         }
 
         self.update_parsers = {key: value for key_tuple, value in self.update_parsers.items() for key in key_tuple}
diff --git a/pyrogram/client/handlers/__init__.py b/pyrogram/client/handlers/__init__.py
@@ -21,10 +21,11 @@
 from .disconnect_handler import DisconnectHandler
 from .inline_query_handler import InlineQueryHandler
 from .message_handler import MessageHandler
+from .poll_handler import PollHandler
 from .raw_update_handler import RawUpdateHandler
 from .user_status_handler import UserStatusHandler
 
 __all__ = [
     "MessageHandler", "DeletedMessagesHandler", "CallbackQueryHandler", "RawUpdateHandler", "DisconnectHandler",
-    "UserStatusHandler", "InlineQueryHandler"
+    "UserStatusHandler", "InlineQueryHandler", "PollHandler"
 ]
diff --git a/pyrogram/client/handlers/poll_handler.py b/pyrogram/client/handlers/poll_handler.py
@@ -0,0 +1,48 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2019 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+from .handler import Handler
+
+
+class PollHandler(Handler):
+    """The Poll handler class. Used to handle polls updates.
+
+    It is intended to be used with :meth:`add_handler() <pyrogram.Client.add_handler>`
+
+    For a nicer way to register this handler, have a look at the
+    :meth:`on_poll() <pyrogram.Client.on_poll>` decorator.
+
+    Args:
+        callback (``callable``):
+            Pass a function that will be called when a new poll update arrives. It takes *(client, poll)*
+            as positional arguments (look at the section below for a detailed description).
+
+        filters (:obj:`Filters <pyrogram.Filters>`):
+            Pass one or more filters to allow only a subset of polls to be passed
+            in your callback function.
+
+    Other parameters:
+        client (:obj:`Client <pyrogram.Client>`):
+            The Client itself, useful when you want to call other API methods inside the poll handler.
+
+        poll (:obj:`Poll <pyrogram.Poll>`):
+            The received poll.
+    """
+
+    def __init__(self, callback: callable, filters=None):
+        super().__init__(callback, filters)
diff --git a/pyrogram/client/methods/decorators/__init__.py b/pyrogram/client/methods/decorators/__init__.py
@@ -21,6 +21,7 @@
 from .on_disconnect import OnDisconnect
 from .on_inline_query import OnInlineQuery
 from .on_message import OnMessage
+from .on_poll import OnPoll
 from .on_raw_update import OnRawUpdate
 from .on_user_status import OnUserStatus
 
@@ -32,6 +33,7 @@ class Decorators(
     OnRawUpdate,
     OnDisconnect,
     OnUserStatus,
-    OnInlineQuery
+    OnInlineQuery,
+    OnPoll
 ):
     pass
diff --git a/pyrogram/client/methods/decorators/on_poll.py b/pyrogram/client/methods/decorators/on_poll.py
@@ -0,0 +1,59 @@
+# Pyrogram - Telegram MTProto API Client Library for Python
+# Copyright (C) 2017-2019 Dan Tès <https://github.com/delivrance>
+#
+# This file is part of Pyrogram.
+#
+# Pyrogram is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Pyrogram is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+from typing import Tuple
+
+import pyrogram
+from pyrogram.client.filters.filter import Filter
+from pyrogram.client.handlers.handler import Handler
+from ...ext import BaseClient
+
+
+class OnPoll(BaseClient):
+    def on_poll(
+        self=None,
+        filters=None,
+        group: int = 0
+    ) -> callable:
+        """Use this decorator to automatically register a function for handling poll updates.
+        This does the same thing as :meth:`add_handler` using the :class:`PollHandler`.
+
+        Args:
+            filters (:obj:`Filters <pyrogram.Filters>`):
+                Pass one or more filters to allow only a subset of polls to be passed
+                in your function.
+
+            group (``int``, *optional*):
+                The group identifier, defaults to 0.
+        """
+
+        def decorator(func: callable) -> Tuple[Handler, int]:
+            if isinstance(func, tuple):
+                func = func[0].callback
+
+            handler = pyrogram.PollHandler(func, filters)
+
+            if isinstance(self, Filter):
+                return pyrogram.PollHandler(func, self), group if filters is None else filters
+
+            if self is not None:
+                self.add_handler(handler, group)
+
+            return handler, group
+
+        return decorator
diff --git a/pyrogram/client/types/messages_and_media/poll.py b/pyrogram/client/types/messages_and_media/poll.py
@@ -22,9 +22,10 @@
 from pyrogram.api import types
 from .poll_option import PollOption
 from ..pyrogram_type import PyrogramType
+from ..update import Update
 
 
-class Poll(PyrogramType):
+class Poll(PyrogramType, Update):
     """This object represents a Poll.
 
     Args:
