Add get_input_* methods for Media and such

telethon/utils.py

@@ -10,8 +10,17 @@ from .tl.types import (
     ChatPhoto, InputPeerChannel, InputPeerChat, InputPeerUser, InputPeerEmpty,
     MessageMediaDocument, MessageMediaPhoto, PeerChannel, InputChannel,
     UserEmpty, InputUser, InputUserEmpty, InputUserSelf, InputPeerSelf,
-    PeerChat, PeerUser, User, UserFull, UserProfilePhoto, Document
-)
+    PeerChat, PeerUser, User, UserFull, UserProfilePhoto, Document,
+    MessageMediaContact, MessageMediaEmpty, MessageMediaGame, MessageMediaGeo,
+    MessageMediaUnsupported, MessageMediaVenue, InputMediaContact,
+    InputMediaDocument, InputMediaEmpty, InputMediaGame,
+    InputMediaGeoPoint, InputMediaPhoto, InputMediaVenue, InputDocument,
+    DocumentEmpty, InputDocumentEmpty, Message, GeoPoint, InputGeoPoint,
+    GeoPointEmpty, InputGeoPointEmpty, Photo, InputPhoto, PhotoEmpty,
+    InputPhotoEmpty, FileLocation, ChatPhotoEmpty, UserProfilePhotoEmpty,
+    FileLocationUnavailable, InputMediaUploadedDocument,
+    InputMediaUploadedPhoto,
+    DocumentAttributeFilename)
 
 
 def get_display_name(entity):
@@ -156,6 +165,161 @@ def get_input_user(entity):
     _raise_cast_fail(entity, 'InputUser')
 
 
+def get_input_document(document):
+    """Similar to get_input_peer, but for documents"""
+    if document is None:
+        return None
+
+    if not isinstance(document, TLObject):
+        _raise_cast_fail(document, 'InputDocument')
+
+    if type(document).subclass_of_id == 0xf33fdb68:  # crc32(b'InputDocument')
+        return document
+
+    if isinstance(document, Document):
+        return InputDocument(id=document.id, access_hash=document.access_hash)
+
+    if isinstance(document, DocumentEmpty):
+        return InputDocumentEmpty()
+
+    if isinstance(document, MessageMediaDocument):
+        return get_input_document(document.document)
+
+    if isinstance(document, Message):
+        return get_input_document(document.media)
+
+    _raise_cast_fail(document, 'InputDocument')
+
+
+def get_input_photo(photo):
+    """Similar to get_input_peer, but for documents"""
+    if photo is None:
+        return None
+
+    if not isinstance(photo, TLObject):
+        _raise_cast_fail(photo, 'InputPhoto')
+
+    if type(photo).subclass_of_id == 0x846363e0:  # crc32(b'InputPhoto')
+        return photo
+
+    if isinstance(photo, Photo):
+        return InputPhoto(id=photo.id, access_hash=photo.access_hash)
+
+    if isinstance(photo, PhotoEmpty):
+        return InputPhotoEmpty()
+
+    _raise_cast_fail(photo, 'InputPhoto')
+
+
+def get_input_geo(geo):
+    """Similar to get_input_peer, but for geo points"""
+    if geo is None:
+        return None
+
+    if not isinstance(geo, TLObject):
+        _raise_cast_fail(geo, 'InputGeoPoint')
+
+    if type(geo).subclass_of_id == 0x430d225:  # crc32(b'InputGeoPoint')
+        return geo
+
+    if isinstance(geo, GeoPoint):
+        return InputGeoPoint(lat=geo.lat, long=geo.long)
+
+    if isinstance(geo, GeoPointEmpty):
+        return InputGeoPointEmpty()
+
+    if isinstance(geo, MessageMediaGeo):
+        return get_input_geo(geo.geo)
+
+    if isinstance(geo, Message):
+        return get_input_geo(geo.media)
+
+    _raise_cast_fail(geo, 'InputGeoPoint')
+
+
+def get_input_media(media, user_caption=None, is_photo=False):
+    """Similar to get_input_peer, but for media.
+
+       If the media is a file location and is_photo is known to be True,
+       it will be treated as an InputMediaUploadedPhoto.
+    """
+    if media is None:
+        return None
+
+    if not isinstance(media, TLObject):
+        _raise_cast_fail(media, 'InputMedia')
+
+    if type(media).subclass_of_id == 0xfaf846f4:  # crc32(b'InputMedia')
+        return media
+
+    if isinstance(media, MessageMediaPhoto):
+        return InputMediaPhoto(
+            id=get_input_photo(media.photo),
+            caption=media.caption if user_caption is None else user_caption,
+            ttl_seconds=media.ttl_seconds
+        )
+
+    if isinstance(media, MessageMediaDocument):
+        return InputMediaDocument(
+            id=get_input_document(media.document),
+            caption=media.caption if user_caption is None else user_caption,
+            ttl_seconds=media.ttl_seconds
+        )
+
+    if isinstance(media, FileLocation):
+        if is_photo:
+            return InputMediaUploadedPhoto(
+                file=media,
+                caption=user_caption or ''
+            )
+        else:
+            return InputMediaUploadedDocument(
+                file=media,
+                mime_type='application/octet-stream',  # unknown, assume bytes
+                attributes=[DocumentAttributeFilename('unnamed')],
+                caption=user_caption or ''
+            )
+
+    if isinstance(media, MessageMediaGame):
+        return InputMediaGame(id=media.game.id)
+
+    if isinstance(media, ChatPhoto) or isinstance(media, UserProfilePhoto):
+        if isinstance(media.photo_big, FileLocationUnavailable):
+            return get_input_media(media.photo_small, is_photo=True)
+        else:
+            return get_input_media(media.photo_big, is_photo=True)
+
+    if isinstance(media, MessageMediaContact):
+        return InputMediaContact(
+            phone_number=media.phone_number,
+            first_name=media.first_name,
+            last_name=media.last_name
+        )
+
+    if isinstance(media, MessageMediaGeo):
+        return InputMediaGeoPoint(geo_point=get_input_geo(media.geo))
+
+    if isinstance(media, MessageMediaVenue):
+        return InputMediaVenue(
+            geo_point=get_input_geo(media.geo),
+            title=media.title,
+            address=media.address,
+            provider=media.provider,
+            venue_id=media.venue_id
+        )
+
+    if any(isinstance(media, t) for t in (
+            MessageMediaEmpty, MessageMediaUnsupported,
+            FileLocationUnavailable, ChatPhotoEmpty,
+            UserProfilePhotoEmpty)):
+        return InputMediaEmpty()
+
+    if isinstance(media, Message):
+        return get_input_media(media.media)
+
+    _raise_cast_fail(media, 'InputMedia')
+
+
 def find_user_or_chat(peer, users, chats):
     """Finds the corresponding user or chat given a peer.
        Returns None if it was not found"""