1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
|
from __future__ import annotations
from typing import Callable
from mypy.plugin import FunctionSigContext, Plugin
from mypy.types import CallableType
class FunctionSigPlugin(Plugin):
def get_function_signature_hook(
self, fullname: str
) -> Callable[[FunctionSigContext], CallableType] | None:
if fullname == "__main__.dynamic_signature":
return my_hook
return None
def my_hook(ctx: FunctionSigContext) -> CallableType:
arg1_args = ctx.args[0]
if len(arg1_args) != 1:
return ctx.default_signature
arg1_type = ctx.api.get_expression_type(arg1_args[0])
return ctx.default_signature.copy_modified(arg_types=[arg1_type], ret_type=arg1_type)
def plugin(version: str) -> type[FunctionSigPlugin]:
return FunctionSigPlugin
|
