Implemented read code on TLObjects Generator

The code generated by the generator now also writes the files
on_response(...) method. Also, all the generated files are
saved in a dictionary containing `constructorId: class`

README.md

@@ -24,5 +24,6 @@ In order to make sure that all the generated files will work, please make sure t
 // vector#1cb5c415 {t:Type} # [ t ] = Vector t;
 ```
 
+Also please make sure to rename `updates#74ae4240 ...` to `updates_tg#74ae4240 ...` or similar to avoid confusion between the updates folder and the updates.py file!
 
  

scheme.tl

@@ -403,7 +403,7 @@ updateShortMessage#914fbf11 flags:# out:flags.1?true mentioned:flags.4?true medi
 updateShortChatMessage#16812688 flags:# out:flags.1?true mentioned:flags.4?true media_unread:flags.5?true silent:flags.13?true id:int from_id:int chat_id:int message:string pts:int pts_count:int date:int fwd_from:flags.2?MessageFwdHeader via_bot_id:flags.11?int reply_to_msg_id:flags.3?int entities:flags.7?Vector<MessageEntity> = Updates;
 updateShort#78d4dec1 update:Update date:int = Updates;
 updatesCombined#725b04c3 updates:Vector<Update> users:Vector<User> chats:Vector<Chat> date:int seq_start:int seq:int = Updates;
-updates#74ae4240 updates:Vector<Update> users:Vector<User> chats:Vector<Chat> date:int seq:int = Updates;
+updates_tg#74ae4240 updates:Vector<Update> users:Vector<User> chats:Vector<Chat> date:int seq:int = Updates;
 updateShortSentMessage#11f1331c flags:# out:flags.1?true id:int pts:int pts_count:int date:int media:flags.9?MessageMedia entities:flags.7?Vector<MessageEntity> = Updates;
 
 photos.photos#8dca6aa5 photos:Vector<Photo> users:Vector<User> = photos.Photos;

tl/tlobject.py

@@ -94,7 +94,10 @@ class TLArg:
         :param generic_definition: Is the argument a generic definition?
                                    (i.e. {X:Type})
         """
-        self.name = name
+        if name == 'self':  # This very only name is restricted
+            self.name = 'is_self'
+        else:
+            self.name = name
 
         # Default values
         self.is_vector = False

tlobjects_generator.py

@@ -10,8 +10,9 @@ def generate_tlobjecs():
     # First ensure that the required parent directories exist
     os.makedirs('tl/functions', exist_ok=True)
     os.makedirs('tl/types', exist_ok=True)
-    for tlobject in TLParser.parse_file('scheme.tl'):
 
+    tlobjects = tuple(TLParser.parse_file('scheme.tl'))
+    for tlobject in tlobjects:
         # Determine the output directory and create it
         out_dir = os.path.join('tl',
                                'functions' if tlobject.is_function
@@ -28,9 +29,8 @@ def generate_tlobjecs():
             open(init_py, 'a').close()
 
         # Create the file
-        filename = os.path.join(out_dir, get_file_name(tlobject))
+        filename = os.path.join(out_dir, get_file_name(tlobject, add_extension=True))
         with open(filename, 'w', encoding='utf-8') as file:
-
             # Let's build the source code!
             with SourceBuilder(file) as builder:
                 builder.writeln('from requests.mtproto_request import MTProtoRequest')
@@ -79,6 +79,10 @@ def generate_tlobjecs():
                     builder.writeln('"""')
 
                 builder.writeln('super().__init__()')
+                # Functions have a result object
+                if tlobject.is_function:
+                    builder.writeln('self.result = None')
+
                 # Leave an empty line if there are any args
                 if args:
                     builder.writeln()
@@ -96,6 +100,45 @@ def generate_tlobjecs():
                     write_onsend_code(builder, arg, tlobject.args)
                 builder.end_block()
 
+                # Write the on_response(self, reader) function
+                builder.writeln('def on_response(self, reader):')
+                # Do not read constructor's ID, since that's already been read somewhere else
+                if tlobject.is_function:
+                    builder.writeln('self.result = reader.tgread_object()')
+                else:
+                    if tlobject.args:
+                        for arg in tlobject.args:
+                            write_onresponse_code(builder, arg, tlobject.args)
+                    else:
+                        builder.writeln('pass')
+                builder.end_block()
+
+    # Once all the objects have been generated, we can now group them in a single file
+    filename = os.path.join('tl', 'all_tlobjects.py')
+    with open(filename, 'w', encoding='utf-8') as file:
+        with SourceBuilder(file) as builder:
+            builder.writeln('"""File generated by TLObjects\' generator. All changes will be ERASED"""')
+            builder.writeln()
+
+            # First add imports
+            for tlobject in tlobjects:
+                builder.writeln('import {}'.format(get_full_file_name(tlobject)))
+            builder.writeln()
+
+            # Then create the dictionary containing constructor_id: class
+            builder.writeln('tlobjects = {')
+            builder.current_indent += 1
+
+            for tlobject in tlobjects:
+                builder.writeln('{}: {}.{},'.format(
+                    hex(tlobject.id),
+                    get_full_file_name(tlobject),
+                    get_class_name(tlobject)
+                ))
+
+            builder.current_indent -= 1
+            builder.writeln('}')
+
 
 def get_class_name(tlobject):
     # Courtesy of http://stackoverflow.com/a/31531797/4759433
@@ -104,17 +147,30 @@ def get_class_name(tlobject):
     return result[:1].upper() + result[1:].replace('_', '')  # Replace again to fully ensure!
 
 
-def get_file_name(tlobject):
+def get_full_file_name(tlobject):
+    fullname = get_file_name(tlobject, add_extension=False)
+    if tlobject.namespace is not None:
+        fullname = '{}.{}'.format(tlobject.namespace, fullname)
+
+    if tlobject.is_function:
+        return 'tl.functions.{}'.format(fullname)
+    else:
+        return 'tl.types.{}'.format(fullname)
+
+
+def get_file_name(tlobject, add_extension):
     # Courtesy of http://stackoverflow.com/a/1176023/4759433
     s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', tlobject.name)
-    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower() + '.py'
+    result = re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()
+    if add_extension:
+        return result + '.py'
+    else:
+        return result
 
 
-foundEver = set()
 def write_onsend_code(builder, arg, args, name=None):
     """
     Writes the write code for the given argument
-
     :param builder: The source code builder
     :param arg: The argument to write
     :param args: All the other arguments in TLObject same on_send. This is required to determine the flags value
@@ -183,9 +239,6 @@ def write_onsend_code(builder, arg, args, name=None):
     else:
         # Else it may be a custom type
         builder.writeln('{}.write(writer)'.format(name))
-        if arg.type not in foundEver:
-            foundEver.add(arg.type)
-            print('{}: {}'.format(arg.type, arg))
 
     # End vector and flag blocks if required (if we opened them before)
     if arg.is_vector:
@@ -195,23 +248,81 @@ def write_onsend_code(builder, arg, args, name=None):
         builder.end_block()
 
 
-''' SourceBuilder generated file example:
+def write_onresponse_code(builder, arg, args, name=None):
+    """
+    Writes the receive code for the given argument
 
-class Example(MTProtoRequest):
-    def __init__(self, some, parameter):
-        """
-        .tl definition: Example#12345678 some:int parameter:int = Exmpl
-        :param some: [type=Vector<int>] Cannot be NONE
-        :param parameter: [type=int] Cannot be NONE
-        """
+    :param builder: The source code builder
+    :param arg: The argument to write
+    :param args: All the other arguments in TLObject same on_send. This is required to determine the flags value
+    :param name: The name of the argument. Defaults to «self.argname»
+                 This argument is an option because it's required when writing Vectors<>
+    """
 
-    def on_send(self, writer):
-        writer.write_int(0x62d6b459)  # example's constructor ID
-        writer.write_int(0x1cb5c415)  # vector code
-        writer.write_int(len(self.msgs))
-        for some_item in self.some:
-            writer.write_int(some_item)
+    if arg.generic_definition:
+        return  # Do nothing, this only specifies a later type
 
-    def on_response(self, reader):
-        pass
-'''
+    if name is None:
+        name = 'self.{}'.format(arg.name)
+
+    # The argument may be a flag, only write that flag was given!
+    if arg.is_flag:
+        builder.writeln('if (flags & (1 << {})) != 0:'.format(arg.flag_index))
+
+    if arg.is_vector:
+        builder.writeln("reader.read_int()  # Vector's constructor ID")
+        builder.writeln('{} = []  # Initialize an empty list'.format(name))
+        builder.writeln('{}_len = reader.read_int()'.format(name))
+        builder.writeln('for _ in range({}_len):'.format(name))
+        # Temporary disable .is_vector, not to enter this if again
+        arg.is_vector = False
+        write_onresponse_code(builder, arg, args, name='{}_item'.format(arg.name))
+        builder.writeln('{}.append({}_item)'.format(name, arg.name))
+        arg.is_vector = True
+
+    elif arg.flag_indicator:
+        # Read the flags, which will indicate what items we should read next
+        builder.writeln('flags = reader.read_int()')
+        builder.writeln()
+
+    elif 'int' == arg.type:
+        builder.writeln('{} = reader.read_int()'.format(name))
+
+    elif 'long' == arg.type:
+        builder.writeln('{} = reader.read_long()'.format(name))
+
+    elif 'int128' == arg.type:
+        builder.writeln('{} = reader.read_large_int(bits=128)'.format(name))
+
+    elif 'int256' == arg.type:
+        builder.writeln('{} = reader.read_large_int(bits=256)'.format(name))
+
+    elif 'double' == arg.type:
+        builder.writeln('{} = reader.read_double()'.format(name))
+
+    elif 'string' == arg.type:
+        builder.writeln('{} = reader.tgread_string()'.format(name))
+
+    elif 'Bool' == arg.type:
+        builder.writeln('{} = reader.tgread_bool()'.format(name))
+
+    elif 'true' == arg.type:  # Awkwardly enough, Telegram has both bool and "true", used in flags
+        builder.writeln('{} = reader.read_int() == 0x3fedd339  # true'.format(name))
+
+    elif 'bytes' == arg.type:
+        builder.writeln('{} = reader.read()'.format(name))
+
+    else:
+        # Else it may be a custom type
+        builder.writeln('{} = reader.tgread_object(reader)'.format(name))
+
+    # End vector and flag blocks if required (if we opened them before)
+    if arg.is_vector:
+        builder.end_block()
+
+    if arg.is_flag:
+        builder.end_block()
+
+
+def get_code(tg_type, stream_name, arg_name):
+    function_name = 'write' if stream_name == 'writer' else 'read'

utils/binary_reader.py

@@ -1,4 +1,5 @@
 from io import BytesIO, BufferedReader
+from tl.all_tlobjects import tlobjects
 import os
 
 
@@ -25,6 +26,9 @@ class BinaryReader:
     def read_long(self, signed=True):
         return int.from_bytes(self.reader.read(8), signed=signed, byteorder='big')
 
+    def read_large_int(self, bits):
+        return int.from_bytes(self.reader.read(bits // 8), byteorder='big')
+
     def read(self, length):
         return self.reader.read(length)
 
@@ -54,6 +58,19 @@ class BinaryReader:
     def tgread_string(self):
         return str(self.tgread_bytes(), encoding='utf-8')
 
+    def tgread_object(self):
+        """Reads a Telegram object"""
+        id = self.read_int()
+        clazz = tlobjects.get(id, None)
+        if clazz is None:
+            raise ImportError('Could not find a matching ID for the TLObject that was supposed to be read. '
+                              'Found ID: {}'.format(hex(id)))
+
+        # Instantiate the class and return the result
+        result = clazz()
+        result.on_response(self)
+        return result
+
     # endregion
 
     def close(self):