| from __future__ import annotations | |
| from typing import Callable | |
| from mypy.nodes import NameExpr | |
| from mypy.plugin import FunctionContext, Plugin | |
| from mypy.types import Instance, NoneType, Type, UnionType, get_proper_type | |
| class AttrPlugin(Plugin): | |
| def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None: | |
| if fullname.startswith("mod.Attr"): | |
| return attr_hook | |
| return None | |
| def attr_hook(ctx: FunctionContext) -> Type: | |
| default = get_proper_type(ctx.default_return_type) | |
| assert isinstance(default, Instance) | |
| if default.type.fullname == "mod.Attr": | |
| attr_base = default | |
| else: | |
| attr_base = None | |
| for base in default.type.bases: | |
| if base.type.fullname == "mod.Attr": | |
| attr_base = base | |
| break | |
| assert attr_base is not None | |
| last_arg_exprs = ctx.args[-1] | |
| if any(isinstance(expr, NameExpr) and expr.name == "True" for expr in last_arg_exprs): | |
| return attr_base | |
| assert len(attr_base.args) == 1 | |
| arg_type = attr_base.args[0] | |
| return Instance( | |
| attr_base.type, | |
| [UnionType([arg_type, NoneType()])], | |
| line=default.line, | |
| column=default.column, | |
| ) | |
| def plugin(version: str) -> type[AttrPlugin]: | |
| return AttrPlugin |