From fce5cfea0e73804081cc25e4b604fadbf4e5f45f Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 2 May 2019 10:19:15 +0200
Subject: [PATCH] Update code generator to emit type hints on init methods

---
 telethon_generator/generators/tlobject.py     | 23 ++++++++-----------
 telethon_generator/parsers/tlobject/parser.py |  2 +-
 telethon_generator/parsers/tlobject/tlarg.py  |  2 +-
 3 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/telethon_generator/generators/tlobject.py b/telethon_generator/generators/tlobject.py
index 64ac3b54..d6b5c4b8 100644
--- a/telethon_generator/generators/tlobject.py
+++ b/telethon_generator/generators/tlobject.py
@@ -89,6 +89,9 @@ def _write_modules(
             # Import struct for the .__bytes__(self) serialization
             builder.writeln('import struct')
 
+            # Import datetime for type hinting
+            builder.writeln('from datetime import datetime')
+
             tlobjects.sort(key=lambda x: x.name)
 
             type_names = set()
@@ -194,24 +197,17 @@ def _write_class_init(tlobject, kind, type_constructors, builder):
     builder.writeln()
 
     # Convert the args to string parameters, flags having =None
-    args = [(a.name if not a.is_flag and not a.can_be_inferred
-             else '{}=None'.format(a.name)) for a in tlobject.real_args]
+    args = ['{}: {}{}'.format(
+        a.name, a.type_hint(), '=None' if a.is_flag or a.can_be_inferred else '')
+        for a in tlobject.real_args
+    ]
 
     # Write the __init__ function if it has any argument
     if not tlobject.real_args:
         return
 
-    builder.writeln('def __init__({}):', ', '.join(['self'] + args))
-    # Write the docstring, to know the type of the args
+    builder.writeln("def __init__({}):", ', '.join(['self'] + args))
     builder.writeln('"""')
-    for arg in tlobject.real_args:
-        if not arg.flag_indicator:
-            builder.writeln(':param {} {}:', arg.type_hint(), arg.name)
-            builder.current_indent -= 1  # It will auto-indent (':')
-
-    # We also want to know what type this request returns
-    # or to which type this constructor belongs to
-    builder.writeln()
     if tlobject.is_function:
         builder.write(':returns {}: ', tlobject.result)
     else:
@@ -232,8 +228,7 @@ def _write_class_init(tlobject, kind, type_constructors, builder):
     # Set the arguments
     for arg in tlobject.real_args:
         if not arg.can_be_inferred:
-            builder.writeln('self.{0} = {0}  # type: {1}',
-                            arg.name, arg.type_hint())
+            builder.writeln('self.{0} = {0}', arg.name)
 
         # Currently the only argument that can be
         # inferred are those called 'random_id'
diff --git a/telethon_generator/parsers/tlobject/parser.py b/telethon_generator/parsers/tlobject/parser.py
index b3caff2a..a502f623 100644
--- a/telethon_generator/parsers/tlobject/parser.py
+++ b/telethon_generator/parsers/tlobject/parser.py
@@ -97,7 +97,7 @@ def parse_tl(file_path, layer, methods=None, ignored_ids=CORE_TYPES):
             if not line:
                 continue
 
-            match = re.match('---(\w+)---', line)
+            match = re.match(r'---(\w+)---', line)
             if match:
                 following_types = match.group(1)
                 is_function = following_types == 'functions'
diff --git a/telethon_generator/parsers/tlobject/tlarg.py b/telethon_generator/parsers/tlobject/tlarg.py
index d34acf25..5dcb3743 100644
--- a/telethon_generator/parsers/tlobject/tlarg.py
+++ b/telethon_generator/parsers/tlobject/tlarg.py
@@ -174,7 +174,7 @@ class TLArg:
             'date': 'Optional[datetime]',  # None date = 0 timestamp
             'bytes': 'bytes',
             'true': 'bool',
-        }.get(cls, "Type{}".format(cls))
+        }.get(cls, "'Type{}'".format(cls))
         if self.is_vector:
             result = 'List[{}]'.format(result)
         if self.is_flag and cls != 'date':