From 3b420b32697af78fc0d1633147c4dbd23288b049 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Damien=20Nad=C3=A9?= Date: Sun, 10 Dec 2023 17:04:16 +0100 Subject: [PATCH] add 'wrapped' type hint information --- tools/generate-tests | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tools/generate-tests b/tools/generate-tests index 950558e..a7bd840 100755 --- a/tools/generate-tests +++ b/tools/generate-tests @@ -52,8 +52,18 @@ class Function: ''' ... ''' return 'return_value'""" - def type(self): - return f"Callable[..., {self.signature.return_type}]" + def type_hint(self, force_coro=False): + if self.signature.params_style == "param": + params = f"[{", ".join(self.signature.params_type for _ in range(self.signature.params_format.count(":")))}]" + elif self.signature.params_style == "kwargs": + params = f"[KwArg({self.signature.params_type})]" + else: + params = f"[VarArg({self.signature.params_type})]" + if self.coroutine or force_coro: + rtype = f"Awaitable[{self.signature.return_type}]" + else: + rtype = f"{self.signature.return_type}" + return f"Callable[{params}, {rtype}]" @property def test_name(self): @@ -137,9 +147,10 @@ class Decoration: return f"""@{self.build} {self.function}""" else: + type_hint = self.function.type_hint(self.decorator_name == "aioretry") return f"""{self.function} -{self.decoration} = {self.build}({self.function.name})""" +{self.decoration}: {type_hint} = {self.build}({self.function.name})""" def decorations(input_data) -> str: @@ -185,7 +196,8 @@ reveal_type(func) import asyncio from typing import Any - +from collections.abc import Callable, Awaitable +from mypy_extensions import VarArg, KwArg from kaioretry import retry, aioretry, Retry, Context