Make possible to attach photos, videos, stickers and locations to polls (#306)

Author: MarsBatya

compiler/docs/compiler.py

@@ -918,6 +918,7 @@ def get_title_list(s: str) -> list:
             InputMediaDocument
             InputMediaPhoto
             InputMediaVideo
+            InputMediaSticker
             InputMessageContent
             InputPhoneContact
             InputPollOption

pyrogram/methods/messages/edit_message_media.py

@@ -17,16 +17,10 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 from datetime import datetime
-import io
-import os
-import re
-from typing import Union
+from typing import Optional, Union
 
 import pyrogram
-from pyrogram import raw
-from pyrogram import types
-from pyrogram import utils
-from pyrogram.file_id import FileType
+from pyrogram import raw, types, utils
 
 
 class EditMessageMedia:
@@ -35,11 +29,10 @@ async def edit_message_media(
         chat_id: Union[int, str],
         message_id: int,
         media: "types.InputMedia",
-        show_caption_above_media: bool = None,
-        schedule_date: datetime = None,
-        business_connection_id: str = None,
-        reply_markup: "types.InlineKeyboardMarkup" = None,
-        file_name: str = None
+        show_caption_above_media: Optional[bool] = None,
+        schedule_date: Optional[datetime] = None,
+        business_connection_id: Optional[str] = None,
+        reply_markup: Optional["types.InlineKeyboardMarkup"] = None,
     ) -> "types.Message":
         """Edit animation, audio, document, photo or video messages, or to add media to text messages.
 
@@ -107,202 +100,38 @@ async def edit_message_media(
         message, entities = None, None
 
         if caption is not None:
-            message, entities = (await utils.parse_text_entities(self, caption, parse_mode, caption_entities)).values()
-
-        if isinstance(media, types.InputMediaPhoto):
-            if isinstance(media.media, io.BytesIO) or os.path.isfile(media.media):
-                uploaded_media = await self.invoke(
-                    raw.functions.messages.UploadMedia(
-                        peer=await self.resolve_peer(chat_id),
-                        media=raw.types.InputMediaUploadedPhoto(
-                            file=await self.save_file(media.media),
-                            spoiler=media.has_spoiler
-                        )
-                    )
-                )
-
-                media = raw.types.InputMediaPhoto(
-                    id=raw.types.InputPhoto(
-                        id=uploaded_media.photo.id,
-                        access_hash=uploaded_media.photo.access_hash,
-                        file_reference=uploaded_media.photo.file_reference
-                    ),
-                    spoiler=media.has_spoiler
-                )
-            elif re.match("^https?://", media.media):
-                media = raw.types.InputMediaPhotoExternal(
-                    url=media.media,
-                    spoiler=media.has_spoiler
-                )
-            else:
-                media = utils.get_input_media_from_file_id(media.media, FileType.PHOTO, has_spoiler=media.has_spoiler)
-        elif isinstance(media, types.InputMediaVideo):
-            if isinstance(media.media, io.BytesIO) or os.path.isfile(media.media):
-                uploaded_media = await self.invoke(
-                    raw.functions.messages.UploadMedia(
-                        peer=await self.resolve_peer(chat_id),
-                        media=raw.types.InputMediaUploadedDocument(
-                            mime_type=self.guess_mime_type(media.media) or "video/mp4",
-                            thumb=await self.save_file(media.thumb),
-                            spoiler=media.has_spoiler,
-                            file=await self.save_file(media.media),
-                            attributes=[
-                                raw.types.DocumentAttributeVideo(
-                                    supports_streaming=media.supports_streaming or None,
-                                    duration=media.duration,
-                                    w=media.width,
-                                    h=media.height
-                                ),
-                                raw.types.DocumentAttributeFilename(
-                                    file_name=file_name or os.path.basename(media.media)
-                                )
-                            ]
-                        )
-                    )
-                )
-
-                media = raw.types.InputMediaDocument(
-                    id=raw.types.InputDocument(
-                        id=uploaded_media.document.id,
-                        access_hash=uploaded_media.document.access_hash,
-                        file_reference=uploaded_media.document.file_reference
-                    ),
-                    spoiler=media.has_spoiler
-                )
-            elif re.match("^https?://", media.media):
-                media = raw.types.InputMediaDocumentExternal(
-                    url=media.media,
-                    spoiler=media.has_spoiler
-                )
-            else:
-                media = utils.get_input_media_from_file_id(media.media, FileType.VIDEO, has_spoiler=media.has_spoiler)
-        elif isinstance(media, types.InputMediaAudio):
-            if isinstance(media.media, io.BytesIO) or os.path.isfile(media.media):
-                media = await self.invoke(
-                    raw.functions.messages.UploadMedia(
-                        peer=await self.resolve_peer(chat_id),
-                        media=raw.types.InputMediaUploadedDocument(
-                            mime_type=self.guess_mime_type(media.media) or "audio/mpeg",
-                            thumb=await self.save_file(media.thumb),
-                            file=await self.save_file(media.media),
-                            attributes=[
-                                raw.types.DocumentAttributeAudio(
-                                    duration=media.duration,
-                                    performer=media.performer,
-                                    title=media.title
-                                ),
-                                raw.types.DocumentAttributeFilename(
-                                    file_name=file_name or os.path.basename(media.media)
-                                )
-                            ]
-                        )
-                    )
-                )
-
-                media = raw.types.InputMediaDocument(
-                    id=raw.types.InputDocument(
-                        id=media.document.id,
-                        access_hash=media.document.access_hash,
-                        file_reference=media.document.file_reference
-                    )
-                )
-            elif re.match("^https?://", media.media):
-                media = raw.types.InputMediaDocumentExternal(
-                    url=media.media
-                )
-            else:
-                media = utils.get_input_media_from_file_id(media.media, FileType.AUDIO)
-        elif isinstance(media, types.InputMediaAnimation):
-            if isinstance(media.media, io.BytesIO) or os.path.isfile(media.media):
-                uploaded_media = await self.invoke(
-                    raw.functions.messages.UploadMedia(
-                        peer=await self.resolve_peer(chat_id),
-                        media=raw.types.InputMediaUploadedDocument(
-                            mime_type=self.guess_mime_type(media.media) or "video/mp4",
-                            thumb=await self.save_file(media.thumb),
-                            spoiler=media.has_spoiler,
-                            file=await self.save_file(media.media),
-                            attributes=[
-                                raw.types.DocumentAttributeVideo(
-                                    supports_streaming=True,
-                                    duration=media.duration,
-                                    w=media.width,
-                                    h=media.height
-                                ),
-                                raw.types.DocumentAttributeFilename(
-                                    file_name=file_name or os.path.basename(media.media)
-                                ),
-                                raw.types.DocumentAttributeAnimated()
-                            ]
-                        )
-                    )
-                )
-
-                media = raw.types.InputMediaDocument(
-                    id=raw.types.InputDocument(
-                        id=uploaded_media.document.id,
-                        access_hash=uploaded_media.document.access_hash,
-                        file_reference=uploaded_media.document.file_reference
-                    ),
-                    spoiler=media.has_spoiler
-                )
-            elif re.match("^https?://", media.media):
-                media = raw.types.InputMediaDocumentExternal(
-                    url=media.media,
-                    spoiler=media.has_spoiler
-                )
-            else:
-                media = utils.get_input_media_from_file_id(media.media, FileType.ANIMATION, has_spoiler=media.has_spoiler)
-        elif isinstance(media, types.InputMediaDocument):
-            if isinstance(media.media, io.BytesIO) or os.path.isfile(media.media):
-                media = await self.invoke(
-                    raw.functions.messages.UploadMedia(
-                        peer=await self.resolve_peer(chat_id),
-                        media=raw.types.InputMediaUploadedDocument(
-                            mime_type=self.guess_mime_type(media.media) or "application/zip",
-                            thumb=await self.save_file(media.thumb),
-                            file=await self.save_file(media.media),
-                            attributes=[
-                                raw.types.DocumentAttributeFilename(
-                                    file_name=file_name or os.path.basename(media.media)
-                                )
-                            ]
-                        )
-                    )
-                )
-
-                media = raw.types.InputMediaDocument(
-                    id=raw.types.InputDocument(
-                        id=media.document.id,
-                        access_hash=media.document.access_hash,
-                        file_reference=media.document.file_reference
-                    )
-                )
-            elif re.match("^https?://", media.media):
-                media = raw.types.InputMediaDocumentExternal(
-                    url=media.media
-                )
-            else:
-                media = utils.get_input_media_from_file_id(media.media, FileType.DOCUMENT)
+            message, entities = (
+                await utils.parse_text_entities(self, caption, parse_mode, caption_entities)
+            ).values()
+
+        if not isinstance(
+            media,
+            (
+                types.InputMediaPhoto,
+                types.InputMediaVideo,
+                types.InputMediaAudio,
+                types.InputMediaAnimation,
+                types.InputMediaDocument,
+            ),
+        ):
+            raise ValueError(f"Unsupported media type {type(media)}")
 
         r = await self.invoke(
             raw.functions.messages.EditMessage(
                 peer=await self.resolve_peer(chat_id),
                 id=message_id,
                 invert_media=show_caption_above_media,
-                media=media,
+                media=await media.write(client=self),
                 schedule_date=utils.datetime_to_timestamp(schedule_date),
                 reply_markup=await reply_markup.write(self) if reply_markup else None,
                 message=message,
-                entities=entities
+                entities=entities,
             ),
-            business_connection_id=business_connection_id
+            business_connection_id=business_connection_id,
         )
 
         for i in r.updates:
             if isinstance(i, (raw.types.UpdateEditMessage, raw.types.UpdateEditChannelMessage)):
                 return await types.Message._parse(
-                    self, i.message,
-                    {i.id: i for i in r.users},
-                    {i.id: i for i in r.chats}
+                    self, i.message, {i.id: i for i in r.users}, {i.id: i for i in r.chats}
                 )

pyrogram/methods/messages/send_poll.py

@@ -31,6 +31,7 @@ async def send_poll(
         chat_id: Union[int, str],
         question: Union[str, "types.FormattedText"],
         options: List[Union[str, "types.InputPollOption"]],
+        media: Optional["types.InputMedia"] = None,
         message_thread_id: Optional[int] = None,
         business_connection_id: Optional[str] = None,
         is_anonymous: bool = True,
@@ -85,6 +86,10 @@ async def send_poll(
             options (List of :obj:`~pyrogram.types.InputPollOption`):
                 List of 2-12 answer options.
 
+            media (:obj:`~pyrogram.types.InputMediaPhoto` | :obj:`~pyrogram.types.InputMediaVideo` | :obj:`~pyrogram.types.InputMediaSticker` | :obj:`~pyrogram.types.Location`, *optional*):
+                Media attached to the poll.
+                Currently supports only photo, video, sticker or location.
+
             message_thread_id (``int``, *optional*):
                 Unique identifier for the target message thread (topic) of the forum.
                 For supergroups only.
@@ -179,6 +184,7 @@ async def send_poll(
         Example:
             .. code-block:: python
 
+                # Regular poll
                 await app.send_poll(
                     chat_id=chat_id,
                     question="Is this a poll question?",
@@ -188,7 +194,38 @@ async def send_poll(
                         types.InputPollOption(text="Maybe")
                     ]
                 )
+
+                # Poll with media
+                await app.send_poll(
+                    chat_id=chat_id,
+                    question="Where we are?",
+                    media=types.InputMediaPhoto("photo.jpg"),
+                    options=[
+                        types.InputPollOption(
+                            text="Maybe here?",
+                            media=types.InputMediaPhoto("photo1.jpg")
+                        ),
+                        types.InputPollOption(
+                            text="Or here?",
+                            media=types.Location(
+                                longitude=49.807760,
+                                latitude=73.088504
+                            ),
+                        ),
+                    ]
+                )
         """
+        if media is not None and not isinstance(
+            media,
+            (
+                types.InputMediaPhoto,
+                types.InputMediaVideo,
+                types.InputMediaSticker,
+                types.Location,
+            ),
+        ):
+            raise ValueError("Unsupported media type")
+
         if isinstance(question, str):
             question = types.FormattedText(text=question)
 
@@ -246,6 +283,7 @@ async def send_poll(
                         close_date=utils.datetime_to_timestamp(close_date)
                     ),
                     correct_answers=correct_option_ids,
+                    attached_media=await media.write(client=self) if media is not None else None,
                     solution=solution,
                     solution_entities=solution_entities
                 ),

pyrogram/types/input_content/__init__.py

@@ -35,6 +35,7 @@
 from .input_media_document import InputMediaDocument
 from .input_media_photo import InputMediaPhoto
 from .input_media_video import InputMediaVideo
+from .input_media_sticker import InputMediaSticker
 from .input_message_content import InputMessageContent
 from .input_phone_contact import InputPhoneContact
 from .input_poll_option import InputPollOption
@@ -77,6 +78,7 @@
     "InputMediaDocument",
     "InputMediaPhoto",
     "InputMediaVideo",
+    "InputMediaSticker",
     "InputMessageContent",
     "InputPhoneContact",
     "InputPollOption",

pyrogram/types/input_content/input_media.py

@@ -16,11 +16,14 @@
 #  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, BinaryIO
+from typing import TYPE_CHECKING, BinaryIO, List, Optional, Union
 
 from ..messages_and_media import MessageEntity
 from ..object import Object
 
+if TYPE_CHECKING:
+    from pyrogram import raw
+
 
 class InputMedia(Object):
     """Content of a media message to be sent.
@@ -32,18 +35,22 @@ class InputMedia(Object):
     - :obj:`~pyrogram.types.InputMediaAudio`
     - :obj:`~pyrogram.types.InputMediaPhoto`
     - :obj:`~pyrogram.types.InputMediaVideo`
+    - :obj:`~pyrogram.types.InputMediaSticker`
     """
 
     def __init__(
         self,
         media: Union[str, BinaryIO],
         caption: str = "",
-        parse_mode: str = None,
-        caption_entities: List[MessageEntity] = None
+        parse_mode: Optional[str] = None,
+        caption_entities: Optional[List[MessageEntity]] = None,
     ):
         super().__init__()
 
         self.media = media
         self.caption = caption
         self.parse_mode = parse_mode
         self.caption_entities = caption_entities
+
+    async def write(self, **kwargs) -> "raw.base.InputMedia":
+        raise NotImplementedError