test-data/unit/plugins/common_api_incremental.py

from __future__ import annotations

from typing import Callable

from mypy.nodes import GDEF, MDEF, Block, ClassDef, SymbolTable, SymbolTableNode, TypeInfo, Var
from mypy.plugin import ClassDefContext, DynamicClassDefContext, Plugin


class DynPlugin(Plugin):
    def get_dynamic_class_hook(
        self, fullname: str
    ) -> Callable[[DynamicClassDefContext], None] | None:
        if fullname == "lib.declarative_base":
            return add_info_hook
        return None

    def get_base_class_hook(self, fullname: str) -> Callable[[ClassDefContext], None] | None:
        sym = self.lookup_fully_qualified(fullname)
        if sym and isinstance(sym.node, TypeInfo):
            if sym.node.metadata.get("magic"):
                return add_magic_hook
        return None


def add_info_hook(ctx: DynamicClassDefContext) -> None:
    class_def = ClassDef(ctx.name, Block([]))
    class_def.fullname = ctx.api.qualified_name(ctx.name)

    info = TypeInfo(SymbolTable(), class_def, ctx.api.cur_mod_id)
    class_def.info = info
    obj = ctx.api.named_type("builtins.object", [])
    info.mro = [info, obj.type]
    info.bases = [obj]
    ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
    info.metadata["magic"] = {"value": True}


def add_magic_hook(ctx: ClassDefContext) -> None:
    info = ctx.cls.info
    str_type = ctx.api.named_type_or_none("builtins.str", [])
    assert str_type is not None
    var = Var("__magic__", str_type)
    var.info = info
    info.names["__magic__"] = SymbolTableNode(MDEF, var)


def plugin(version: str) -> type[DynPlugin]:
    return DynPlugin