attrs field plugin

gistfile1.txt

from typing import (
    Callable,
    Optional,
    Type,
)

from mypy.plugin import FunctionContext, Plugin
from mypy.nodes import (
    Block,
    ClassDef,
    MDEF,
    SymbolTableNode,
    Var,
    TypeInfo,
    SymbolTable,
)
from mypy.types import Instance


def make_var(n: str, i: Instance) -> Var:
    res = Var(n, i)
    res.info = i.type
    return res


def adjust_fields(fc: FunctionContext) -> Type:
    attrs_class = fc.arg_types[0][0].ret_type.type
    attribute_i: Instance = fc.default_return_type.type.bases[0].args[0]
    fields = [
        (k, v.type) for k, v in attrs_class.names.items() if isinstance(v.node, Var)
    ]
    sym_table = SymbolTable(
        {
            n: SymbolTableNode(MDEF, make_var(n, attribute_i.copy_modified(args=[t])))
            for (n, t) in fields
        }
    )
    cd = ClassDef(f"{attrs_class.name}Attributes", Block([]))
    cd.fullname = f"{attrs_class.name}Attributes"
    ti = TypeInfo(
        sym_table,
        cd,
        "attr",
    )
    ti.mro = [ti, attrs_class.mro[-1]]
    ti.type_vars = []
    res = Instance(
        ti,
        [],
    )
    return res


class CustomPlugin(Plugin):
    def get_function_hook(
        self, fullname: str
    ) -> Optional[Callable[[FunctionContext], Type]]:
        if fullname == "attr.fields":
            return adjust_fields

        return None


def plugin(_: str):
    return CustomPlugin