diff --git a/telethon_generator/tl_generator.py b/telethon_generator/tl_generator.py index a524ca7b..4838087c 100644 --- a/telethon_generator/tl_generator.py +++ b/telethon_generator/tl_generator.py @@ -6,6 +6,7 @@ from zlib import crc32 from collections import defaultdict from .parser import SourceBuilder, TLParser, TLObject + AUTO_GEN_NOTICE = \ '"""File generated by TLObjects\' generator. All changes will be ERASED"""' @@ -140,6 +141,12 @@ class TLGenerator: 'from {}.tl.tlobject import TLObject'.format('.' * depth) ) + # No need to import base in base.py + if ns: + builder.writeln( + 'from . import base' + ) + # Import 'get_input_*' utils # TODO Support them on types too if 'functions' in out_dir: @@ -343,7 +350,7 @@ class TLGenerator: builder.writeln('def from_reader(reader):') for arg in tlobject.args: TLGenerator.write_read_code( - builder, arg, tlobject.args, name='_' + arg.name + builder, arg, tlobject, name='_' + arg.name ) builder.writeln('return {}({})'.format( @@ -550,15 +557,14 @@ class TLGenerator: return True # Something was written @staticmethod - def write_read_code(builder, arg, args, name): + def write_read_code(builder, arg, tlobject, name): """ Writes the read code for the given argument, setting the arg.name variable to its read value. :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 tlobject: The TLObject for which from_reader() method is written :param name: The name of the argument. Defaults to "self.argname" This argument is an option because it's required when writing Vectors<> @@ -595,7 +601,7 @@ class TLGenerator: builder.writeln('for _ in range(reader.read_int()):') # Temporary disable .is_vector, not to enter this if again arg.is_vector = False - TLGenerator.write_read_code(builder, arg, args, name='_x') + TLGenerator.write_read_code(builder, arg, tlobject, name='_x') builder.writeln('{}.append(_x)'.format(name)) arg.is_vector = True @@ -644,9 +650,8 @@ class TLGenerator: if not arg.skip_constructor_id: builder.writeln('{} = reader.tgread_object()'.format(name)) else: - builder.writeln('from . import {}'.format(TLObject.class_name_for(arg.type))) - builder.writeln('{} = {}.from_reader(reader)'.format( - name, TLObject.class_name_for(arg.type))) + builder.writeln('{} = {}{}.from_reader(reader)'.format( + name, 'base.' if tlobject.namespace else '', TLObject.class_name_for(arg.type))) # End vector and flag blocks if required (if we opened them before) if arg.is_vector: