| 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 |
