| from mypy.plugin import Plugin |
| from mypy.nodes import NameExpr |
| from mypy.types import UnionType, NoneType, Instance |
| |
| class AttrPlugin(Plugin): |
| def get_function_hook(self, fullname): |
| if fullname.startswith('mod.Attr'): |
| return attr_hook |
| return None |
| |
| def attr_hook(ctx): |
| assert isinstance(ctx.default_return_type, Instance) |
| if ctx.default_return_type.type.fullname == 'mod.Attr': |
| attr_base = ctx.default_return_type |
| else: |
| attr_base = None |
| for base in ctx.default_return_type.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=ctx.default_return_type.line, |
| column=ctx.default_return_type.column) |
| |
| def plugin(version): |
| return AttrPlugin |
