codekingpro/portable-devtools

1# ext/mypy/plugin.py
2# Copyright (C) 2021-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7
8"""
9Mypy plugin for SQLAlchemy ORM.
10
11"""
12from __future__ import annotations
13
14from typing import Callable
15from typing import List
16from typing import Optional
17from typing import Tuple
18from typing import Type as TypingType
19from typing import Union
20
21from mypy import nodes
22from mypy.mro import calculate_mro
23from mypy.mro import MroError
24from mypy.nodes import Block
25from mypy.nodes import ClassDef
26from mypy.nodes import GDEF
27from mypy.nodes import MypyFile
28from mypy.nodes import NameExpr
29from mypy.nodes import SymbolTable
30from mypy.nodes import SymbolTableNode
31from mypy.nodes import TypeInfo
32from mypy.plugin import AttributeContext
33from mypy.plugin import ClassDefContext
34from mypy.plugin import DynamicClassDefContext
35from mypy.plugin import Plugin
36from mypy.plugin import SemanticAnalyzerPluginInterface
37from mypy.types import get_proper_type
38from mypy.types import Instance
39from mypy.types import Type
40
41from . import decl_class
42from . import names
43from . import util
44
45try:
46    __import__("sqlalchemy-stubs")
47except ImportError:
48    pass
49else:
50    raise ImportError(
51        "The SQLAlchemy mypy plugin in SQLAlchemy "
52        "2.0 does not work with sqlalchemy-stubs or "
53        "sqlalchemy2-stubs installed, as well as with any other third party "
54        "SQLAlchemy stubs.  Please uninstall all SQLAlchemy stubs "
55        "packages."
56    )
57
58
59class SQLAlchemyPlugin(Plugin):
60    def get_dynamic_class_hook(
61        self, fullname: str
62    ) -> Optional[Callable[[DynamicClassDefContext], None]]:
63        if names.type_id_for_fullname(fullname) is names.DECLARATIVE_BASE:
64            return _dynamic_class_hook
65        return None
66
67    def get_customize_class_mro_hook(
68        self, fullname: str
69    ) -> Optional[Callable[[ClassDefContext], None]]:
70        return _fill_in_decorators
71
72    def get_class_decorator_hook(
73        self, fullname: str
74    ) -> Optional[Callable[[ClassDefContext], None]]:
75        sym = self.lookup_fully_qualified(fullname)
76
77        if sym is not None and sym.node is not None:
78            type_id = names.type_id_for_named_node(sym.node)
79            if type_id is names.MAPPED_DECORATOR:
80                return _cls_decorator_hook
81            elif type_id in (
82                names.AS_DECLARATIVE,
83                names.AS_DECLARATIVE_BASE,
84            ):
85                return _base_cls_decorator_hook
86            elif type_id is names.DECLARATIVE_MIXIN:
87                return _declarative_mixin_hook
88
89        return None
90
91    def get_metaclass_hook(
92        self, fullname: str
93    ) -> Optional[Callable[[ClassDefContext], None]]:
94        if names.type_id_for_fullname(fullname) is names.DECLARATIVE_META:
95            # Set any classes that explicitly have metaclass=DeclarativeMeta
96            # as declarative so the check in `get_base_class_hook()` works
97            return _metaclass_cls_hook
98
99        return None
100
101    def get_base_class_hook(
102        self, fullname: str
103    ) -> Optional[Callable[[ClassDefContext], None]]:
104        sym = self.lookup_fully_qualified(fullname)
105
106        if (
107            sym
108            and isinstance(sym.node, TypeInfo)
109            and util.has_declarative_base(sym.node)
110        ):
111            return _base_cls_hook
112
113        return None
114
115    def get_attribute_hook(
116        self, fullname: str
117    ) -> Optional[Callable[[AttributeContext], Type]]:
118        if fullname.startswith(
119            "sqlalchemy.orm.attributes.QueryableAttribute."
120        ):
121            return _queryable_getattr_hook
122
123        return None
124
125    def get_additional_deps(
126        self, file: MypyFile
127    ) -> List[Tuple[int, str, int]]:
128        return [
129            #
130            (10, "sqlalchemy.orm", -1),
131            (10, "sqlalchemy.orm.attributes", -1),
132            (10, "sqlalchemy.orm.decl_api", -1),
133        ]
134
135
136def plugin(version: str) -> TypingType[SQLAlchemyPlugin]:
137    return SQLAlchemyPlugin
138
139
140def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None:
141    """Generate a declarative Base class when the declarative_base() function
142    is encountered."""
143
144    _add_globals(ctx)
145
146    cls = ClassDef(ctx.name, Block([]))
147    cls.fullname = ctx.api.qualified_name(ctx.name)
148
149    info = TypeInfo(SymbolTable(), cls, ctx.api.cur_mod_id)
150    cls.info = info
151    _set_declarative_metaclass(ctx.api, cls)
152
153    cls_arg = util.get_callexpr_kwarg(ctx.call, "cls", expr_types=(NameExpr,))
154    if cls_arg is not None and isinstance(cls_arg.node, TypeInfo):
155        util.set_is_base(cls_arg.node)
156        decl_class.scan_declarative_assignments_and_apply_types(
157            cls_arg.node.defn, ctx.api, is_mixin_scan=True
158        )
159        info.bases = [Instance(cls_arg.node, [])]
160    else:
161        obj = ctx.api.named_type(names.NAMED_TYPE_BUILTINS_OBJECT)
162
163        info.bases = [obj]
164
165    try:
166        calculate_mro(info)
167    except MroError:
168        util.fail(
169            ctx.api, "Not able to calculate MRO for declarative base", ctx.call
170        )
171        obj = ctx.api.named_type(names.NAMED_TYPE_BUILTINS_OBJECT)
172        info.bases = [obj]
173        info.fallback_to_any = True
174
175    ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
176    util.set_is_base(info)
177
178
179def _fill_in_decorators(ctx: ClassDefContext) -> None:
180    for decorator in ctx.cls.decorators:
181        # set the ".fullname" attribute of a class decorator
182        # that is a MemberExpr.   This causes the logic in
183        # semanal.py->apply_class_plugin_hooks to invoke the
184        # get_class_decorator_hook for our "registry.map_class()"
185        # and "registry.as_declarative_base()" methods.
186        # this seems like a bug in mypy that these decorators are otherwise
187        # skipped.
188
189        if (
190            isinstance(decorator, nodes.CallExpr)
191            and isinstance(decorator.callee, nodes.MemberExpr)
192            and decorator.callee.name == "as_declarative_base"
193        ):
194            target = decorator.callee
195        elif (
196            isinstance(decorator, nodes.MemberExpr)
197            and decorator.name == "mapped"
198        ):
199            target = decorator
200        else:
201            continue
202
203        if isinstance(target.expr, NameExpr):
204            sym = ctx.api.lookup_qualified(
205                target.expr.name, target, suppress_errors=True
206            )
207        else:
208            continue
209
210        if sym and sym.node:
211            sym_type = get_proper_type(sym.type)
212            if isinstance(sym_type, Instance):
213                target.fullname = f"{sym_type.type.fullname}.{target.name}"
214            else:
215                # if the registry is in the same file as where the
216                # decorator is used, it might not have semantic
217                # symbols applied and we can't get a fully qualified
218                # name or an inferred type, so we are actually going to
219                # flag an error in this case that they need to annotate
220                # it.  The "registry" is declared just
221                # once (or few times), so they have to just not use
222                # type inference for its assignment in this one case.
223                util.fail(
224                    ctx.api,
225                    "Class decorator called %s(), but we can't "
226                    "tell if it's from an ORM registry.  Please "
227                    "annotate the registry assignment, e.g. "
228                    "my_registry: registry = registry()" % target.name,
229                    sym.node,
230                )
231
232
233def _cls_decorator_hook(ctx: ClassDefContext) -> None:
234    _add_globals(ctx)
235    assert isinstance(ctx.reason, nodes.MemberExpr)
236    expr = ctx.reason.expr
237
238    assert isinstance(expr, nodes.RefExpr) and isinstance(expr.node, nodes.Var)
239
240    node_type = get_proper_type(expr.node.type)
241
242    assert (
243        isinstance(node_type, Instance)
244        and names.type_id_for_named_node(node_type.type) is names.REGISTRY
245    )
246
247    decl_class.scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
248
249
250def _base_cls_decorator_hook(ctx: ClassDefContext) -> None:
251    _add_globals(ctx)
252
253    cls = ctx.cls
254
255    _set_declarative_metaclass(ctx.api, cls)
256
257    util.set_is_base(ctx.cls.info)
258    decl_class.scan_declarative_assignments_and_apply_types(
259        cls, ctx.api, is_mixin_scan=True
260    )
261
262
263def _declarative_mixin_hook(ctx: ClassDefContext) -> None:
264    _add_globals(ctx)
265    util.set_is_base(ctx.cls.info)
266    decl_class.scan_declarative_assignments_and_apply_types(
267        ctx.cls, ctx.api, is_mixin_scan=True
268    )
269
270
271def _metaclass_cls_hook(ctx: ClassDefContext) -> None:
272    util.set_is_base(ctx.cls.info)
273
274
275def _base_cls_hook(ctx: ClassDefContext) -> None:
276    _add_globals(ctx)
277    decl_class.scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
278
279
280def _queryable_getattr_hook(ctx: AttributeContext) -> Type:
281    # how do I....tell it it has no attribute of a certain name?
282    # can't find any Type that seems to match that
283    return ctx.default_attr_type
284
285
286def _add_globals(ctx: Union[ClassDefContext, DynamicClassDefContext]) -> None:
287    """Add __sa_DeclarativeMeta and __sa_Mapped symbol to the global space
288    for all class defs
289
290    """
291
292    util.add_global(ctx, "sqlalchemy.orm", "Mapped", "__sa_Mapped")
293
294
295def _set_declarative_metaclass(
296    api: SemanticAnalyzerPluginInterface, target_cls: ClassDef
297) -> None:
298    info = target_cls.info
299    sym = api.lookup_fully_qualified_or_none(
300        "sqlalchemy.orm.decl_api.DeclarativeMeta"
301    )
302    assert sym is not None and isinstance(sym.node, TypeInfo)
303    info.declared_metaclass = info.metaclass_type = Instance(sym.node, [])
304