Add support for public polls

KurimuzonAkuma · Oct 21, 2024 · 046312a · 046312a
1 parent 055d125
commit 046312a
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 34 deletions.
diff --git a/pyrogram/dispatcher.py b/pyrogram/dispatcher.py
@@ -22,9 +22,7 @@
 from collections import OrderedDict
 
 import pyrogram
-from pyrogram import errors
 from pyrogram import utils
-from pyrogram import raw
 from pyrogram.handlers import (
     CallbackQueryHandler, MessageHandler, EditedMessageHandler, DeletedMessagesHandler,
     UserStatusHandler, RawUpdateHandler, InlineQueryHandler, PollHandler, PreCheckoutQueryHandler,
@@ -41,7 +39,7 @@
     UpdateBotInlineSend, UpdateChatParticipant, UpdateChannelParticipant,
     UpdateBotChatInviteRequester, UpdateStory, UpdateBotShippingQuery, UpdateBotMessageReaction,
     UpdateBotMessageReactions, UpdateBotChatBoost, UpdateBusinessBotCallbackQuery,
-    UpdateBotPurchasedPaidMedia
+    UpdateBotPurchasedPaidMedia, UpdateMessagePollVote
 )
 
 log = logging.getLogger(__name__)
@@ -55,7 +53,7 @@ class Dispatcher:
     CHAT_MEMBER_UPDATES = (UpdateChatParticipant, UpdateChannelParticipant)
     USER_STATUS_UPDATES = (UpdateUserStatus,)
     BOT_INLINE_QUERY_UPDATES = (UpdateBotInlineQuery,)
-    POLL_UPDATES = (UpdateMessagePoll,)
+    POLL_UPDATES = (UpdateMessagePoll, UpdateMessagePollVote)
     CHOSEN_INLINE_RESULT_UPDATES = (UpdateBotInlineSend,)
     CHAT_JOIN_REQUEST_UPDATES = (UpdateBotChatInviteRequester,)
     NEW_STORY_UPDATES = (UpdateStory,)
@@ -128,7 +126,7 @@ async def inline_query_parser(update, users, chats):
 
         async def poll_parser(update, users, chats):
             return (
-                pyrogram.types.Poll._parse_update(self.client, update),
+                pyrogram.types.Poll._parse_update(self.client, update, users),
                 PollHandler
             )
 

diff --git a/pyrogram/types/messages_and_media/poll.py b/pyrogram/types/messages_and_media/poll.py
@@ -77,6 +77,9 @@ class Poll(Object, Update):
 
         close_date (:py:obj:`~datetime.datetime`, *optional*):
             Point in time when the poll will be automatically closed.
+
+        voter (:obj:`~pyrogram.types.User`, *optional*):
+            The user that voted in the poll.
     """
 
     def __init__(
@@ -97,7 +100,8 @@ def __init__(
         explanation: Optional[str] = None,
         explanation_entities: Optional[List["types.MessageEntity"]] = None,
         open_period: Optional[int] = None,
-        close_date: Optional[datetime] = None
+        close_date: Optional[datetime] = None,
+        voter: Optional["types.User"] = None
     ):
         super().__init__(client)
 
@@ -116,6 +120,7 @@ def __init__(
         self.explanation_entities = explanation_entities
         self.open_period = open_period
         self.close_date = close_date
+        self.voter = voter
 
     @staticmethod
     def _parse(client, media_poll: Union["raw.types.MessageMediaPoll", "raw.types.UpdateMessagePoll"]) -> "Poll":
@@ -195,38 +200,57 @@ def _parse(client, media_poll: Union["raw.types.MessageMediaPoll", "raw.types.Up
         )
 
     @staticmethod
-    def _parse_update(client, update: "raw.types.UpdateMessagePoll"):
-        if update.poll is not None:
-            return Poll._parse(client, update)
+    def _parse_update(client, update: Union["raw.types.UpdateMessagePoll", "raw.types.UpdateMessagePollVote"], users: dict):
+        if isinstance(update, raw.types.UpdateMessagePoll):
+            if update.poll is not None:
+                return Poll._parse(client, update)
 
-        results = update.results.results
-        chosen_option_id = None
-        correct_option_id = None
-        options = []
+            results = update.results.results
+            chosen_option_id = None
+            correct_option_id = None
+            options = []
 
-        for i, result in enumerate(results):
-            if result.chosen:
-                chosen_option_id = i
+            for i, result in enumerate(results):
+                if result.chosen:
+                    chosen_option_id = i
 
-            if result.correct:
-                correct_option_id = i
+                if result.correct:
+                    correct_option_id = i
 
-            options.append(
-                types.PollOption(
-                    text="",
-                    voter_count=result.voters,
-                    data=result.option,
-                    client=client
+                options.append(
+                    types.PollOption(
+                        text="",
+                        voter_count=result.voters,
+                        data=result.option,
+                        client=client
+                    )
                 )
+
+            return Poll(
+                id=str(update.poll_id),
+                question="",
+                options=options,
+                total_voter_count=update.results.total_voters,
+                is_closed=False,
+                chosen_option_id=chosen_option_id,
+                correct_option_id=correct_option_id,
+                client=client
             )
 
-        return Poll(
-            id=str(update.poll_id),
-            question="",
-            options=options,
-            total_voter_count=update.results.total_voters,
-            is_closed=False,
-            chosen_option_id=chosen_option_id,
-            correct_option_id=correct_option_id,
-            client=client
-        )
+        if isinstance(update, raw.types.UpdateMessagePollVote):
+            return Poll(
+                id=str(update.poll_id),
+                question="",
+                options=[
+                    types.PollOption(
+                        text="",
+                        voter_count=None,
+                        data=option,
+                        client=client
+                    ) for option in update.options
+                ],
+                total_voter_count=None,
+                is_closed=False,
+                voter=types.User._parse(client, users[update.peer.user_id]),
+                client=client
+            )