Skip to content

Commit 1368540

Browse files
committed
added generic return type for method invoke
1 parent 173f977 commit 1368540

4 files changed

Lines changed: 36 additions & 9 deletions

File tree

compiler/api/compiler.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,19 @@ def camel(s: str):
9494
return "".join([i[0].upper() + i[1:] for i in s.split("_")])
9595

9696

97+
# noinspection PyShadowingBuiltins, PyShadowingNames
98+
def get_return_type_hint(qualtype: str) -> str:
99+
"""Get return type hint for generic TLObject"""
100+
if qualtype.startswith("Vector"):
101+
# Extract inner type from Vector<Type>
102+
inner = qualtype.split("<")[1][:-1]
103+
ns, name = inner.split(".") if "." in inner else ("", inner)
104+
return f'"List[raw.base.{".".join([ns, name]).strip(".")}]"'
105+
else:
106+
ns, name = qualtype.split(".") if "." in qualtype else ("", qualtype)
107+
return f'"raw.base.{".".join([ns, name]).strip(".")}"'
108+
109+
97110
# noinspection PyShadowingBuiltins, PyShadowingNames
98111
def get_type_hint(type: str) -> str:
99112
is_flag = FLAGS_RE.match(type)
@@ -541,6 +554,12 @@ def start(format: bool = False):
541554
slots = ", ".join([f'"{i[0]}"' for i in sorted_args])
542555
return_arguments = ", ".join([f"{i[0]}={i[0]}" for i in sorted_args])
543556

557+
# Generate generic type hint for functions
558+
if c.section == "functions":
559+
generic_type = f"[{get_return_type_hint(c.qualtype)}]"
560+
else:
561+
generic_type = ""
562+
544563
compiled_combinator = combinator_tmpl.format(
545564
notice=notice,
546565
warning=WARNING,
@@ -553,7 +572,8 @@ def start(format: bool = False):
553572
fields=fields,
554573
read_types=read_types,
555574
write_types=write_types,
556-
return_arguments=return_arguments
575+
return_arguments=return_arguments,
576+
generic_type=generic_type
557577
)
558578

559579
directory = "types" if c.section == "types" else c.section

compiler/api/template/combinator.txt

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,18 @@
11
{notice}
22

33
from io import BytesIO
4+
from typing import TYPE_CHECKING, List, Optional, Any
45

56
from pyrogram.raw.core.primitives import Int, Long, Int128, Int256, Bool, Bytes, String, Double, Vector
67
from pyrogram.raw.core import TLObject
7-
from pyrogram import raw
8-
from typing import List, Optional, Any
8+
9+
if TYPE_CHECKING:
10+
from pyrogram import raw
911

1012
{warning}
1113

1214

13-
class {name}(TLObject): # type: ignore
15+
class {name}(TLObject{generic_type}):
1416
"""{docstring}
1517
"""
1618

pyrogram/methods/advanced/invoke.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
# along with Pyrogram. If not, see <http://www.gnu.org/licenses/>.
1818

1919
import logging
20+
from typing import TypeVar, overload
2021

2122
import pyrogram
2223
from pyrogram import raw
@@ -25,18 +26,20 @@
2526

2627
log = logging.getLogger(__name__)
2728

29+
ReturnType = TypeVar('ReturnType')
30+
2831

2932
class Invoke:
3033
async def invoke(
3134
self: "pyrogram.Client",
32-
query: TLObject,
35+
query: TLObject[ReturnType],
3336
retries: int = Session.MAX_RETRIES,
3437
timeout: float = Session.WAIT_TIMEOUT,
3538
sleep_threshold: float = None,
3639
retry_delay: float = Session.RETRY_DELAY,
3740
recaptcha_token: str = None,
3841
business_connection_id: str = None
39-
):
42+
) -> ReturnType:
4043
"""Invoke raw Telegram functions.
4144
4245
This method makes it possible to manually call every single Telegram API method in a low-level manner.

pyrogram/raw/core/tl_object.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,14 @@
1818

1919
from io import BytesIO
2020
from json import dumps
21-
from typing import cast, List, Any, Union, Dict
21+
from typing import cast, List, Any, Union, Dict, TypeVar, Generic
2222

2323
from ..all import objects
2424

25+
ReturnType = TypeVar("ReturnType")
2526

26-
class TLObject:
27+
28+
class TLObject(Generic[ReturnType]):
2729
__slots__: List[str] = []
2830

2931
QUALNAME = "Base"
@@ -78,5 +80,5 @@ def __eq__(self, other: Any) -> bool:
7880
def __len__(self) -> int:
7981
return len(self.write())
8082

81-
def __call__(self, *args: Any, **kwargs: Any) -> Any:
83+
def __call__(self, *args: Any, **kwargs: Any) -> ReturnType:
8284
pass

0 commit comments

Comments
 (0)