codekingpro/portable-devtools

1# ext/mypy/apply.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
8from __future__ import annotations
9
10from typing import List
11from typing import Optional
12from typing import Union
13
14from mypy.nodes import ARG_NAMED_OPT
15from mypy.nodes import Argument
16from mypy.nodes import AssignmentStmt
17from mypy.nodes import CallExpr
18from mypy.nodes import ClassDef
19from mypy.nodes import MDEF
20from mypy.nodes import MemberExpr
21from mypy.nodes import NameExpr
22from mypy.nodes import RefExpr
23from mypy.nodes import StrExpr
24from mypy.nodes import SymbolTableNode
25from mypy.nodes import TempNode
26from mypy.nodes import TypeInfo
27from mypy.nodes import Var
28from mypy.plugin import SemanticAnalyzerPluginInterface
29from mypy.plugins.common import add_method_to_class
30from mypy.types import AnyType
31from mypy.types import get_proper_type
32from mypy.types import Instance
33from mypy.types import NoneTyp
34from mypy.types import ProperType
35from mypy.types import TypeOfAny
36from mypy.types import UnboundType
37from mypy.types import UnionType
38
39from . import infer
40from . import util
41from .names import expr_to_mapped_constructor
42from .names import NAMED_TYPE_SQLA_MAPPED
43
44
45def apply_mypy_mapped_attr(
46    cls: ClassDef,
47    api: SemanticAnalyzerPluginInterface,
48    item: Union[NameExpr, StrExpr],
49    attributes: List[util.SQLAlchemyAttribute],
50) -> None:
51    if isinstance(item, NameExpr):
52        name = item.name
53    elif isinstance(item, StrExpr):
54        name = item.value
55    else:
56        return None
57
58    for stmt in cls.defs.body:
59        if (
60            isinstance(stmt, AssignmentStmt)
61            and isinstance(stmt.lvalues[0], NameExpr)
62            and stmt.lvalues[0].name == name
63        ):
64            break
65    else:
66        util.fail(api, f"Can't find mapped attribute {name}", cls)
67        return None
68
69    if stmt.type is None:
70        util.fail(
71            api,
72            "Statement linked from _mypy_mapped_attrs has no "
73            "typing information",
74            stmt,
75        )
76        return None
77
78    left_hand_explicit_type = get_proper_type(stmt.type)
79    assert isinstance(
80        left_hand_explicit_type, (Instance, UnionType, UnboundType)
81    )
82
83    attributes.append(
84        util.SQLAlchemyAttribute(
85            name=name,
86            line=item.line,
87            column=item.column,
88            typ=left_hand_explicit_type,
89            info=cls.info,
90        )
91    )
92
93    apply_type_to_mapped_statement(
94        api, stmt, stmt.lvalues[0], left_hand_explicit_type, None
95    )
96
97
98def re_apply_declarative_assignments(
99    cls: ClassDef,
100    api: SemanticAnalyzerPluginInterface,
101    attributes: List[util.SQLAlchemyAttribute],
102) -> None:
103    """For multiple class passes, re-apply our left-hand side types as mypy
104    seems to reset them in place.
105
106    """
107    mapped_attr_lookup = {attr.name: attr for attr in attributes}
108    update_cls_metadata = False
109
110    for stmt in cls.defs.body:
111        # for a re-apply, all of our statements are AssignmentStmt;
112        # @declared_attr calls will have been converted and this
113        # currently seems to be preserved by mypy (but who knows if this
114        # will change).
115        if (
116            isinstance(stmt, AssignmentStmt)
117            and isinstance(stmt.lvalues[0], NameExpr)
118            and stmt.lvalues[0].name in mapped_attr_lookup
119            and isinstance(stmt.lvalues[0].node, Var)
120        ):
121            left_node = stmt.lvalues[0].node
122
123            python_type_for_type = mapped_attr_lookup[
124                stmt.lvalues[0].name
125            ].type
126
127            left_node_proper_type = get_proper_type(left_node.type)
128
129            # if we have scanned an UnboundType and now there's a more
130            # specific type than UnboundType, call the re-scan so we
131            # can get that set up correctly
132            if (
133                isinstance(python_type_for_type, UnboundType)
134                and not isinstance(left_node_proper_type, UnboundType)
135                and (
136                    isinstance(stmt.rvalue, CallExpr)
137                    and isinstance(stmt.rvalue.callee, MemberExpr)
138                    and isinstance(stmt.rvalue.callee.expr, NameExpr)
139                    and stmt.rvalue.callee.expr.node is not None
140                    and stmt.rvalue.callee.expr.node.fullname
141                    == NAMED_TYPE_SQLA_MAPPED
142                    and stmt.rvalue.callee.name == "_empty_constructor"
143                    and isinstance(stmt.rvalue.args[0], CallExpr)
144                    and isinstance(stmt.rvalue.args[0].callee, RefExpr)
145                )
146            ):
147                new_python_type_for_type = (
148                    infer.infer_type_from_right_hand_nameexpr(
149                        api,
150                        stmt,
151                        left_node,
152                        left_node_proper_type,
153                        stmt.rvalue.args[0].callee,
154                    )
155                )
156
157                if new_python_type_for_type is not None and not isinstance(
158                    new_python_type_for_type, UnboundType
159                ):
160                    python_type_for_type = new_python_type_for_type
161
162                    # update the SQLAlchemyAttribute with the better
163                    # information
164                    mapped_attr_lookup[stmt.lvalues[0].name].type = (
165                        python_type_for_type
166                    )
167
168                    update_cls_metadata = True
169
170            if (
171                not isinstance(left_node.type, Instance)
172                or left_node.type.type.fullname != NAMED_TYPE_SQLA_MAPPED
173            ):
174                assert python_type_for_type is not None
175                left_node.type = api.named_type(
176                    NAMED_TYPE_SQLA_MAPPED, [python_type_for_type]
177                )
178
179    if update_cls_metadata:
180        util.set_mapped_attributes(cls.info, attributes)
181
182
183def apply_type_to_mapped_statement(
184    api: SemanticAnalyzerPluginInterface,
185    stmt: AssignmentStmt,
186    lvalue: NameExpr,
187    left_hand_explicit_type: Optional[ProperType],
188    python_type_for_type: Optional[ProperType],
189) -> None:
190    """Apply the Mapped[<type>] annotation and right hand object to a
191    declarative assignment statement.
192
193    This converts a Python declarative class statement such as::
194
195        class User(Base):
196            # ...
197
198            attrname = Column(Integer)
199
200    To one that describes the final Python behavior to Mypy::
201
202    ... format: off
203
204        class User(Base):
205            # ...
206
207            attrname : Mapped[Optional[int]] = <meaningless temp node>
208
209    ... format: on
210
211    """
212    left_node = lvalue.node
213    assert isinstance(left_node, Var)
214
215    # to be completely honest I have no idea what the difference between
216    # left_node.type and stmt.type is, what it means if these are different
217    # vs. the same, why in order to get tests to pass I have to assign
218    # to stmt.type for the second case and not the first.  this is complete
219    # trying every combination until it works stuff.
220
221    if left_hand_explicit_type is not None:
222        lvalue.is_inferred_def = False
223        left_node.type = api.named_type(
224            NAMED_TYPE_SQLA_MAPPED, [left_hand_explicit_type]
225        )
226    else:
227        lvalue.is_inferred_def = False
228        left_node.type = api.named_type(
229            NAMED_TYPE_SQLA_MAPPED,
230            (
231                [AnyType(TypeOfAny.special_form)]
232                if python_type_for_type is None
233                else [python_type_for_type]
234            ),
235        )
236
237    # so to have it skip the right side totally, we can do this:
238    # stmt.rvalue = TempNode(AnyType(TypeOfAny.special_form))
239
240    # however, if we instead manufacture a new node that uses the old
241    # one, then we can still get type checking for the call itself,
242    # e.g. the Column, relationship() call, etc.
243
244    # rewrite the node as:
245    # <attr> : Mapped[<typ>] =
246    # _sa_Mapped._empty_constructor(<original CallExpr from rvalue>)
247    # the original right-hand side is maintained so it gets type checked
248    # internally
249    stmt.rvalue = expr_to_mapped_constructor(stmt.rvalue)
250
251    if stmt.type is not None and python_type_for_type is not None:
252        stmt.type = python_type_for_type
253
254
255def add_additional_orm_attributes(
256    cls: ClassDef,
257    api: SemanticAnalyzerPluginInterface,
258    attributes: List[util.SQLAlchemyAttribute],
259) -> None:
260    """Apply __init__, __table__ and other attributes to the mapped class."""
261
262    info = util.info_for_cls(cls, api)
263
264    if info is None:
265        return
266
267    is_base = util.get_is_base(info)
268
269    if "__init__" not in info.names and not is_base:
270        mapped_attr_names = {attr.name: attr.type for attr in attributes}
271
272        for base in info.mro[1:-1]:
273            if "sqlalchemy" not in info.metadata:
274                continue
275
276            base_cls_attributes = util.get_mapped_attributes(base, api)
277            if base_cls_attributes is None:
278                continue
279
280            for attr in base_cls_attributes:
281                mapped_attr_names.setdefault(attr.name, attr.type)
282
283        arguments = []
284        for name, typ in mapped_attr_names.items():
285            if typ is None:
286                typ = AnyType(TypeOfAny.special_form)
287            arguments.append(
288                Argument(
289                    variable=Var(name, typ),
290                    type_annotation=typ,
291                    initializer=TempNode(typ),
292                    kind=ARG_NAMED_OPT,
293                )
294            )
295
296        add_method_to_class(api, cls, "__init__", arguments, NoneTyp())
297
298    if "__table__" not in info.names and util.get_has_table(info):
299        _apply_placeholder_attr_to_class(
300            api, cls, "sqlalchemy.sql.schema.Table", "__table__"
301        )
302    if not is_base:
303        _apply_placeholder_attr_to_class(
304            api, cls, "sqlalchemy.orm.mapper.Mapper", "__mapper__"
305        )
306
307
308def _apply_placeholder_attr_to_class(
309    api: SemanticAnalyzerPluginInterface,
310    cls: ClassDef,
311    qualified_name: str,
312    attrname: str,
313) -> None:
314    sym = api.lookup_fully_qualified_or_none(qualified_name)
315    if sym:
316        assert isinstance(sym.node, TypeInfo)
317        type_: ProperType = Instance(sym.node, [])
318    else:
319        type_ = AnyType(TypeOfAny.special_form)
320    var = Var(attrname)
321    var._fullname = cls.fullname + "." + attrname
322    var.info = cls.info
323    var.type = type_
324    cls.info.names[attrname] = SymbolTableNode(MDEF, var)
325