Update poll parsing

delivrance · delivrance · commit f9aacd814a60 · 2022-09-06T19:04:47.000+02:00
diff --git a/pyrogram/types/messages_and_media/message.py b/pyrogram/types/messages_and_media/message.py
@@ -827,7 +827,8 @@ async def _parse(
                     except MessageIdsEmpty:
                         pass
 
-            client.message_cache[(parsed_message.chat.id, parsed_message.id)] = parsed_message
+            if not parsed_message.poll:  # Do not cache poll messages
+                client.message_cache[(parsed_message.chat.id, parsed_message.id)] = parsed_message
 
             return parsed_message
 
diff --git a/pyrogram/types/messages_and_media/poll.py b/pyrogram/types/messages_and_media/poll.py
@@ -16,10 +16,11 @@
 #  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 List, Union
+from datetime import datetime
+from typing import List, Union, Optional
 
 import pyrogram
-from pyrogram import raw, enums
+from pyrogram import raw, enums, utils
 from pyrogram import types
 from ..object import Object
 from ..update import Update
@@ -53,8 +54,26 @@ class Poll(Object, Update):
         allows_multiple_answers (``bool``, *optional*):
             True, if the poll allows multiple answers.
 
-        chosen_option (``int``, *optional*):
-            Index of your chosen option (0-9), None in case you haven't voted yet.
+        chosen_option_id (``int``, *optional*):
+            0-based index of the chosen option), None in case of no vote yet.
+
+        correct_option_id (``int``, *optional*):
+            0-based identifier of the correct answer option.
+            Available only for polls in the quiz mode, which are closed, or was sent (not forwarded) by the bot or to
+            the private chat with the bot.
+
+        explanation (``str``, *optional*):
+            Text that is shown when a user chooses an incorrect answer or taps on the lamp icon in a quiz-style poll,
+            0-200 characters.
+
+        explanation_entities (List of :obj:`~pyrogram.types.MessageEntity`, *optional*):
+            Special entities like usernames, URLs, bot commands, etc. that appear in the explanation.
+
+        open_period (``int``, *optional*):
+            Amount of time in seconds the poll will be active after creation.
+
+        close_date (:py:obj:`~datetime.datetime`, *optional*):
+            Point in time when the poll will be automatically closed.
     """
 
     def __init__(
@@ -69,8 +88,12 @@ def __init__(
         is_anonymous: bool = None,
         type: "enums.PollType" = None,
         allows_multiple_answers: bool = None,
-        # correct_option_id: int,
-        chosen_option: int = None
+        chosen_option_id: Optional[int] = None,
+        correct_option_id: Optional[int] = None,
+        explanation: Optional[str] = None,
+        explanation_entities: Optional[List["types.MessageEntity"]] = None,
+        open_period: Optional[int] = None,
+        close_date: Optional[datetime] = None
     ):
         super().__init__(client)
 
@@ -82,14 +105,21 @@ def __init__(
         self.is_anonymous = is_anonymous
         self.type = type
         self.allows_multiple_answers = allows_multiple_answers
-        # self.correct_option_id = correct_option_id
-        self.chosen_option = chosen_option
+        self.chosen_option_id = chosen_option_id
+        self.correct_option_id = correct_option_id
+        self.explanation = explanation
+        self.explanation_entities = explanation_entities
+        self.open_period = open_period
+        self.close_date = close_date
 
     @staticmethod
     def _parse(client, media_poll: Union["raw.types.MessageMediaPoll", "raw.types.UpdateMessagePoll"]) -> "Poll":
-        poll = media_poll.poll  # type: raw.types.Poll
-        results = media_poll.results.results
-        chosen_option = None
+        poll: raw.types.Poll = media_poll.poll
+        poll_results: raw.types.PollResults = media_poll.results
+        results: List[raw.types.PollAnswerVoters] = poll_results.results
+
+        chosen_option_id = None
+        correct_option_id = None
         options = []
 
         for i, answer in enumerate(poll.answers):
@@ -100,7 +130,10 @@ def _parse(client, media_poll: Union["raw.types.MessageMediaPoll", "raw.types.Up
                 voter_count = result.voters
 
                 if result.chosen:
-                    chosen_option = i
+                    chosen_option_id = i
+
+                if result.correct:
+                    correct_option_id = i
 
             options.append(
                 types.PollOption(
@@ -120,7 +153,15 @@ def _parse(client, media_poll: Union["raw.types.MessageMediaPoll", "raw.types.Up
             is_anonymous=not poll.public_voters,
             type=enums.PollType.QUIZ if poll.quiz else enums.PollType.REGULAR,
             allows_multiple_answers=poll.multiple_choice,
-            chosen_option=chosen_option,
+            chosen_option_id=chosen_option_id,
+            correct_option_id=correct_option_id,
+            explanation=poll_results.solution,
+            explanation_entities=[
+                types.MessageEntity._parse(client, i, {})
+                for i in poll_results.solution_entities
+            ] if poll_results.solution_entities else None,
+            open_period=poll.close_period,
+            close_date=utils.timestamp_to_datetime(poll.close_date),
             client=client
         )
 
@@ -130,12 +171,16 @@ def _parse_update(client, update: "raw.types.UpdateMessagePoll"):
             return Poll._parse(client, update)
 
         results = update.results.results
-        chosen_option = None
+        chosen_option_id = None
+        correct_option_id = None
         options = []
 
         for i, result in enumerate(results):
             if result.chosen:
-                chosen_option = i
+                chosen_option_id = i
+
+            if result.correct:
+                correct_option_id = i
 
             options.append(
                 types.PollOption(
@@ -152,6 +197,7 @@ def _parse_update(client, update: "raw.types.UpdateMessagePoll"):
             options=options,
             total_voter_count=update.results.total_voters,
             is_closed=False,
-            chosen_option=chosen_option,
+            chosen_option_id=chosen_option_id,
+            correct_option_id=correct_option_id,
             client=client
         )
