Add generic return type for method invoke

KurimuzonAkuma/kurigram#252
Co-authored-by: ZeN220 <[email protected]>

Author: SpEcHiDe

compiler/api/compiler.py

@@ -98,6 +98,19 @@ def camel(s: str):
     return "".join([i[0].upper() + i[1:] for i in s.split("_")])
 
 
+# noinspection PyShadowingBuiltins, PyShadowingNames
+def get_return_type_hint(qualtype: str) -> str:
+    """Get return type hint for generic TLObject"""
+    if qualtype.startswith("Vector"):
+        # Extract inner type from Vector<Type>
+        inner = qualtype.split("<")[1][:-1]
+        ns, name = inner.split(".") if "." in inner else ("", inner)
+        return f'"List[raw.base.{".".join([ns, name]).strip(".")}]"'
+    else:
+        ns, name = qualtype.split(".") if "." in qualtype else ("", qualtype)
+        return f'"raw.base.{".".join([ns, name]).strip(".")}"'
+
+
 # noinspection PyShadowingBuiltins, PyShadowingNames
 def get_type_hint(type: str) -> str:
     is_flag = FLAGS_RE.match(type)
@@ -546,6 +559,12 @@ def start(format: bool = False):
         slots = ", ".join([f'"{i[0]}"' for i in sorted_args])
         return_arguments = ", ".join([f"{i[0]}={i[0]}" for i in sorted_args])
 
+        # Generate generic type hint for functions
+        if c.section == "functions":
+            generic_type = f"[{get_return_type_hint(c.qualtype)}]"
+        else:
+            generic_type = ""
+
         compiled_combinator = combinator_tmpl.format(
             notice=notice,
             warning=WARNING,
@@ -558,7 +577,8 @@ def start(format: bool = False):
             fields=fields,
             read_types=read_types,
             write_types=write_types,
-            return_arguments=return_arguments
+            return_arguments=return_arguments,
+            generic_type=generic_type
         )
 
         directory = "types" if c.section == "types" else c.section

compiler/api/template/combinator.txt

@@ -1,16 +1,18 @@
 {notice}
 
 from io import BytesIO
+from typing import TYPE_CHECKING, Optional, Any
 
 from pyrogram.raw.core.primitives import Int, Long, Int128, Int256, Bool, Bytes, String, Double, Vector
 from pyrogram.raw.core import TLObject
-from pyrogram import raw
-from typing import Optional, Any
+
+if TYPE_CHECKING:
+    from pyrogram import raw
 
 {warning}
 
 
-class {name}(TLObject):  # type: ignore
+class {name}(TLObject{generic_type}):
     """{docstring}
     """
 

pyrogram/methods/advanced/invoke.py

@@ -17,6 +17,7 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 import logging
+from typing import TypeVar
 
 import pyrogram
 from pyrogram import raw
@@ -25,15 +26,17 @@
 
 log = logging.getLogger(__name__)
 
+ReturnType = TypeVar('ReturnType')
+
 
 class Invoke:
     async def invoke(
         self: "pyrogram.Client",
-        query: TLObject,
+        query: TLObject[ReturnType],
         retries: int = Session.MAX_RETRIES,
         timeout: float = Session.WAIT_TIMEOUT,
         sleep_threshold: float = None
-    ):
+    ) -> ReturnType:
         """Invoke raw Telegram functions.
 
         This method makes it possible to manually call every single Telegram API method in a low-level manner.

pyrogram/raw/core/tl_object.py

@@ -18,12 +18,14 @@
 
 from io import BytesIO
 from json import dumps
-from typing import cast, Any, Union
+from typing import cast, Any, Union, TypeVar, Generic
 
 from ..all import objects
 
+ReturnType = TypeVar("ReturnType")
 
-class TLObject:
+
+class TLObject(Generic[ReturnType]):
     __slots__: list[str] = []
 
     QUALNAME = "Base"
@@ -78,5 +80,5 @@ def __eq__(self, other: Any) -> bool:
     def __len__(self) -> int:
         return len(self.write())
 
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+    def __call__(self, *args: Any, **kwargs: Any) -> ReturnType:
         pass