1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
|
from __future__ import annotations
from typing import Callable
from mypy.nodes import GDEF, Block, ClassDef, SymbolTable, SymbolTableNode, TypeInfo, Var
from mypy.plugin import ClassDefContext, DynamicClassDefContext, Plugin
from mypy.types import Instance, get_proper_type
DECL_BASES: set[str] = set()
class DynPlugin(Plugin):
def get_dynamic_class_hook(
self, fullname: str
) -> Callable[[DynamicClassDefContext], None] | None:
if fullname == "mod.declarative_base":
return add_info_hook
return None
def get_base_class_hook(self, fullname: str) -> Callable[[ClassDefContext], None] | None:
if fullname in DECL_BASES:
return replace_col_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))
DECL_BASES.add(class_def.fullname)
def replace_col_hook(ctx: ClassDefContext) -> None:
info = ctx.cls.info
for sym in info.names.values():
node = sym.node
if isinstance(node, Var) and isinstance(
(node_type := get_proper_type(node.type)), Instance
):
if node_type.type.fullname == "mod.Column":
new_sym = ctx.api.lookup_fully_qualified_or_none("mod.Instr")
if new_sym:
new_info = new_sym.node
assert isinstance(new_info, TypeInfo)
node.type = Instance(
new_info, node_type.args, node_type.line, node_type.column
)
def plugin(version: str) -> type[DynPlugin]:
return DynPlugin
|
