Implement automatic cast to InputUser too (closes #159)

telethon/utils.py

@@ -8,6 +8,7 @@ from .tl.types import (
     Channel, ChannelForbidden, Chat, ChatEmpty, ChatForbidden, ChatFull,
     ChatPhoto, InputPeerChannel, InputPeerChat, InputPeerUser,
     MessageMediaDocument, MessageMediaPhoto, PeerChannel, InputChannel,
+    UserEmpty, InputUser, InputUserEmpty, InputUserSelf,
     PeerChat, PeerUser, User, UserFull, UserProfilePhoto)
 
 
@@ -81,8 +82,7 @@ def get_input_peer(entity):
 
 
 def get_input_channel(entity):
-    """Gets the input peer for the given "entity" (user, chat or channel).
-       A ValueError is raised if the given entity isn't a supported type."""
+    """Similar to get_input_peer, but for InputChannel's alone"""
     if type(entity).subclass_of_id == 0x40f202fd:  # crc32(b'InputChannel')
         return entity
 
@@ -93,6 +93,27 @@ def get_input_channel(entity):
                      .format(type(entity).__name__))
 
 
+def get_input_user(entity):
+    """Similar to get_input_peer, but for InputUser's alone"""
+    if type(entity).subclass_of_id == 0xe669bf46:  # crc32(b'InputUser')
+        return entity
+
+    if isinstance(entity, User):
+        if entity.is_self:
+            return InputUserSelf()
+        else:
+            return InputUser(entity.id, entity.access_hash)
+
+    if isinstance(entity, UserEmpty):
+        return InputUserEmpty()
+
+    if isinstance(entity, UserFull):
+        return get_input_user(entity.user)
+
+    raise ValueError('Cannot cast {} to any kind of InputUser.'
+                     .format(type(entity).__name__))
+
+
 def find_user_or_chat(peer, users, chats):
     """Finds the corresponding user or chat given a peer.
        Returns None if it was not found"""

telethon_generator/tl_generator.py

@@ -184,16 +184,20 @@ class TLGenerator:
                         .format('.' * depth))
 
         if tlobject.is_function:
-            if any(a for a in tlobject.args if a.type == 'InputPeer'):
-                # We can automatically convert a normal peer to an InputPeer,
-                # it will make invoking a lot of requests a lot simpler.
-                builder.writeln('from {}.utils import get_input_peer'
-                                .format('.' * depth))
+            util_imports = set()
+            for a in tlobject.args:
+                # We can automatically convert some "full" types to
+                # "input only" (like User -> InputPeerUser, etc.)
+                if a.type == 'InputPeer':
+                    util_imports.add('get_input_peer')
+                elif a.type == 'InputChannel':
+                    util_imports.add('get_input_channel')
+                elif a.type == 'InputUser':
+                    util_imports.add('get_input_user')
 
-            if any(a for a in tlobject.args if a.type == 'InputChannel'):
-                # Same applies to channels
-                builder.writeln('from {}.utils import get_input_channel'
-                                .format('.' * depth))
+            if util_imports:
+                builder.writeln('from {}.utils import {}'.format(
+                    '.' * depth, ', '.join(util_imports)))
 
         if any(a for a in tlobject.args if a.can_be_inferred):
             # Currently only 'random_id' needs 'os' to be imported
@@ -327,6 +331,10 @@ class TLGenerator:
                 builder.writeln(
                     'self.{0} = get_input_channel({0})'.format(arg.name)
                 )
+            elif arg.type == 'InputUser' and tlobject.is_function:
+                builder.writeln(
+                    'self.{0} = get_input_user({0})'.format(arg.name)
+                )
             else:
                 builder.writeln('self.{0} = {0}'.format(arg.name))