From bdee94eaf3332f3c32828c0e70feee81a7084db7 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Mon, 10 Jul 2017 16:04:10 +0200 Subject: [PATCH] Implement automatic cast to InputUser too (closes #159) --- telethon/utils.py | 25 +++++++++++++++++++++++-- telethon_generator/tl_generator.py | 26 +++++++++++++++++--------- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/telethon/utils.py b/telethon/utils.py index 411db593..362fe6bb 100644 --- a/telethon/utils.py +++ b/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""" diff --git a/telethon_generator/tl_generator.py b/telethon_generator/tl_generator.py index 0fffe1b8..e6e3dc4d 100644 --- a/telethon_generator/tl_generator.py +++ b/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))