Add core type hints for generated classes

delivrance · delivrance · commit 7c2c8783336f · 2018-04-25T12:55:38.000+02:00
diff --git a/compiler/api/compiler.py b/compiler/api/compiler.py
@@ -99,6 +99,32 @@ def get_references(t: str):
     return t
 
 
+def get_argument_type(arg):
+    is_flag = FLAGS_RE.match(arg[1])
+    name, t = arg
+
+    if is_flag:
+        t = t.split("?")[1]
+
+    if t in core_types:
+        if t == "long" or "int" in t:
+            t = ": int"
+        elif t == "double":
+            t = ": float"
+        elif t == "string":
+            t = ": str"
+        else:
+            t = ": {}".format(t.lower())
+    elif t == "true":
+        t = ": bool"
+    elif t.startswith("Vector"):
+        t = ": list"
+    else:
+        return name + ("=None" if is_flag else "")
+
+    return name + t + (" = None" if is_flag else "")
+
+
 class Combinator:
     def __init__(self,
                  section: str,
@@ -263,10 +289,7 @@ def start():
         sorted_args = sort_args(c.args)
 
         arguments = ", " + ", ".join(
-            ["{}{}".format(
-                i[0],
-                "=None" if FLAGS_RE.match(i[1]) else ""
-            ) for i in sorted_args]
+            [get_argument_type(i) for i in sorted_args]
         ) if c.args else ""
 
         fields = "\n        ".join(
diff --git a/pyrogram/api/core/__init__.py b/pyrogram/api/core/__init__.py
@@ -22,4 +22,7 @@
 from .message import Message
 from .msg_container import MsgContainer
 from .object import Object
-from .primitives import *
+from .primitives import (
+    Bool, BoolTrue, BoolFalse, Bytes, Double,
+    Int, Long, Int128, Int256, Null, String, Vector
+)
