codekingpro/portable-devtools
115k
1# ext/mypy/decl_class.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 AssignmentStmt
15from mypy.nodes import CallExpr
16from mypy.nodes import ClassDef
17from mypy.nodes import Decorator
18from mypy.nodes import LambdaExpr
19from mypy.nodes import ListExpr
20from mypy.nodes import MemberExpr
21from mypy.nodes import NameExpr
22from mypy.nodes import PlaceholderNode
23from mypy.nodes import RefExpr
24from mypy.nodes import StrExpr
25from mypy.nodes import SymbolNode
26from mypy.nodes import SymbolTableNode
27from mypy.nodes import TempNode
28from mypy.nodes import TypeInfo
29from mypy.nodes import Var
30from mypy.plugin import SemanticAnalyzerPluginInterface
31from mypy.types import AnyType
32from mypy.types import CallableType
33from mypy.types import get_proper_type
34from mypy.types import Instance
35from mypy.types import NoneType
36from mypy.types import ProperType
37from mypy.types import Type
38from mypy.types import TypeOfAny
39from mypy.types import UnboundType
40from mypy.types import UnionType
41
42from . import apply
43from . import infer
44from . import names
45from . import util
46
47
48def scan_declarative_assignments_and_apply_types(
49 cls: ClassDef,
50 api: SemanticAnalyzerPluginInterface,
51 is_mixin_scan: bool = False,
52) -> Optional[List[util.SQLAlchemyAttribute]]:
53 info = util.info_for_cls(cls, api)
54
55 if info is None:
56 # this can occur during cached passes
57 return None
58 elif cls.fullname.startswith("builtins"):
59 return None
60
61 mapped_attributes: Optional[List[util.SQLAlchemyAttribute]] = (
62 util.get_mapped_attributes(info, api)
63 )
64
65 # used by assign.add_additional_orm_attributes among others
66 util.establish_as_sqlalchemy(info)
67
68 if mapped_attributes is not None:
69 # ensure that a class that's mapped is always picked up by
70 # its mapped() decorator or declarative metaclass before
71 # it would be detected as an unmapped mixin class
72
73 if not is_mixin_scan:
74 # mypy can call us more than once. it then *may* have reset the
75 # left hand side of everything, but not the right that we removed,
76 # removing our ability to re-scan. but we have the types
77 # here, so lets re-apply them, or if we have an UnboundType,
78 # we can re-scan
79
80 apply.re_apply_declarative_assignments(cls, api, mapped_attributes)
81
82 return mapped_attributes
83
84 mapped_attributes = []
85
86 if not cls.defs.body:
87 # when we get a mixin class from another file, the body is
88 # empty (!) but the names are in the symbol table. so use that.
89
90 for sym_name, sym in info.names.items():
91 _scan_symbol_table_entry(
92 cls, api, sym_name, sym, mapped_attributes
93 )
94 else:
95 for stmt in util.flatten_typechecking(cls.defs.body):
96 if isinstance(stmt, AssignmentStmt):
97 _scan_declarative_assignment_stmt(
98 cls, api, stmt, mapped_attributes
99 )
100 elif isinstance(stmt, Decorator):
101 _scan_declarative_decorator_stmt(
102 cls, api, stmt, mapped_attributes
103 )
104 _scan_for_mapped_bases(cls, api)
105
106 if not is_mixin_scan:
107 apply.add_additional_orm_attributes(cls, api, mapped_attributes)
108
109 util.set_mapped_attributes(info, mapped_attributes)
110
111 return mapped_attributes
112
113
114def _scan_symbol_table_entry(
115 cls: ClassDef,
116 api: SemanticAnalyzerPluginInterface,
117 name: str,
118 value: SymbolTableNode,
119 attributes: List[util.SQLAlchemyAttribute],
120) -> None:
121 """Extract mapping information from a SymbolTableNode that's in the
122 type.names dictionary.
123
124 """
125 value_type = get_proper_type(value.type)
126 if not isinstance(value_type, Instance):
127 return
128
129 left_hand_explicit_type = None
130 type_id = names.type_id_for_named_node(value_type.type)
131 # type_id = names._type_id_for_unbound_type(value.type.type, cls, api)
132
133 err = False
134
135 # TODO: this is nearly the same logic as that of
136 # _scan_declarative_decorator_stmt, likely can be merged
137 if type_id in {
138 names.MAPPED,
139 names.RELATIONSHIP,
140 names.COMPOSITE_PROPERTY,
141 names.MAPPER_PROPERTY,
142 names.SYNONYM_PROPERTY,
143 names.COLUMN_PROPERTY,
144 }:
145 if value_type.args:
146 left_hand_explicit_type = get_proper_type(value_type.args[0])
147 else:
148 err = True
149 elif type_id is names.COLUMN:
150 if not value_type.args:
151 err = True
152 else:
153 typeengine_arg: Union[ProperType, TypeInfo] = get_proper_type(
154 value_type.args[0]
155 )
156 if isinstance(typeengine_arg, Instance):
157 typeengine_arg = typeengine_arg.type
158
159 if isinstance(typeengine_arg, (UnboundType, TypeInfo)):
160 sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
161 if sym is not None and isinstance(sym.node, TypeInfo):
162 if names.has_base_type_id(sym.node, names.TYPEENGINE):
163 left_hand_explicit_type = UnionType(
164 [
165 infer.extract_python_type_from_typeengine(
166 api, sym.node, []
167 ),
168 NoneType(),
169 ]
170 )
171 else:
172 util.fail(
173 api,
174 "Column type should be a TypeEngine "
175 "subclass not '{}'".format(sym.node.fullname),
176 value_type,
177 )
178
179 if err:
180 msg = (
181 "Can't infer type from attribute {} on class {}. "
182 "please specify a return type from this function that is "
183 "one of: Mapped[<python type>], relationship[<target class>], "
184 "Column[<TypeEngine>], MapperProperty[<python type>]"
185 )
186 util.fail(api, msg.format(name, cls.name), cls)
187
188 left_hand_explicit_type = AnyType(TypeOfAny.special_form)
189
190 if left_hand_explicit_type is not None:
191 assert value.node is not None
192 attributes.append(
193 util.SQLAlchemyAttribute(
194 name=name,
195 line=value.node.line,
196 column=value.node.column,
197 typ=left_hand_explicit_type,
198 info=cls.info,
199 )
200 )
201
202
203def _scan_declarative_decorator_stmt(
204 cls: ClassDef,
205 api: SemanticAnalyzerPluginInterface,
206 stmt: Decorator,
207 attributes: List[util.SQLAlchemyAttribute],
208) -> None:
209 """Extract mapping information from a @declared_attr in a declarative
210 class.
211
212 E.g.::
213
214 @reg.mapped
215 class MyClass:
216 # ...
217
218 @declared_attr
219 def updated_at(cls) -> Column[DateTime]:
220 return Column(DateTime)
221
222 Will resolve in mypy as::
223
224 @reg.mapped
225 class MyClass:
226 # ...
227
228 updated_at: Mapped[Optional[datetime.datetime]]
229
230 """
231 for dec in stmt.decorators:
232 if (
233 isinstance(dec, (NameExpr, MemberExpr, SymbolNode))
234 and names.type_id_for_named_node(dec) is names.DECLARED_ATTR
235 ):
236 break
237 else:
238 return
239
240 dec_index = cls.defs.body.index(stmt)
241
242 left_hand_explicit_type: Optional[ProperType] = None
243
244 if util.name_is_dunder(stmt.name):
245 # for dunder names like __table_args__, __tablename__,
246 # __mapper_args__ etc., rewrite these as simple assignment
247 # statements; otherwise mypy doesn't like if the decorated
248 # function has an annotation like ``cls: Type[Foo]`` because
249 # it isn't @classmethod
250 any_ = AnyType(TypeOfAny.special_form)
251 left_node = NameExpr(stmt.var.name)
252 left_node.node = stmt.var
253 new_stmt = AssignmentStmt([left_node], TempNode(any_))
254 new_stmt.type = left_node.node.type
255 cls.defs.body[dec_index] = new_stmt
256 return
257 elif isinstance(stmt.func.type, CallableType):
258 func_type = stmt.func.type.ret_type
259 if isinstance(func_type, UnboundType):
260 type_id = names.type_id_for_unbound_type(func_type, cls, api)
261 else:
262 # this does not seem to occur unless the type argument is
263 # incorrect
264 return
265
266 if (
267 type_id
268 in {
269 names.MAPPED,
270 names.RELATIONSHIP,
271 names.COMPOSITE_PROPERTY,
272 names.MAPPER_PROPERTY,
273 names.SYNONYM_PROPERTY,
274 names.COLUMN_PROPERTY,
275 }
276 and func_type.args
277 ):
278 left_hand_explicit_type = get_proper_type(func_type.args[0])
279 elif type_id is names.COLUMN and func_type.args:
280 typeengine_arg = func_type.args[0]
281 if isinstance(typeengine_arg, UnboundType):
282 sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
283 if sym is not None and isinstance(sym.node, TypeInfo):
284 if names.has_base_type_id(sym.node, names.TYPEENGINE):
285 left_hand_explicit_type = UnionType(
286 [
287 infer.extract_python_type_from_typeengine(
288 api, sym.node, []
289 ),
290 NoneType(),
291 ]
292 )
293 else:
294 util.fail(
295 api,
296 "Column type should be a TypeEngine "
297 "subclass not '{}'".format(sym.node.fullname),
298 func_type,
299 )
300
301 if left_hand_explicit_type is None:
302 # no type on the decorated function. our option here is to
303 # dig into the function body and get the return type, but they
304 # should just have an annotation.
305 msg = (
306 "Can't infer type from @declared_attr on function '{}'; "
307 "please specify a return type from this function that is "
308 "one of: Mapped[<python type>], relationship[<target class>], "
309 "Column[<TypeEngine>], MapperProperty[<python type>]"
310 )
311 util.fail(api, msg.format(stmt.var.name), stmt)
312
313 left_hand_explicit_type = AnyType(TypeOfAny.special_form)
314
315 left_node = NameExpr(stmt.var.name)
316 left_node.node = stmt.var
317
318 # totally feeling around in the dark here as I don't totally understand
319 # the significance of UnboundType. It seems to be something that is
320 # not going to do what's expected when it is applied as the type of
321 # an AssignmentStatement. So do a feeling-around-in-the-dark version
322 # of converting it to the regular Instance/TypeInfo/UnionType structures
323 # we see everywhere else.
324 if isinstance(left_hand_explicit_type, UnboundType):
325 left_hand_explicit_type = get_proper_type(
326 util.unbound_to_instance(api, left_hand_explicit_type)
327 )
328
329 left_node.node.type = api.named_type(
330 names.NAMED_TYPE_SQLA_MAPPED, [left_hand_explicit_type]
331 )
332
333 # this will ignore the rvalue entirely
334 # rvalue = TempNode(AnyType(TypeOfAny.special_form))
335
336 # rewrite the node as:
337 # <attr> : Mapped[<typ>] =
338 # _sa_Mapped._empty_constructor(lambda: <function body>)
339 # the function body is maintained so it gets type checked internally
340 rvalue = names.expr_to_mapped_constructor(
341 LambdaExpr(stmt.func.arguments, stmt.func.body)
342 )
343
344 new_stmt = AssignmentStmt([left_node], rvalue)
345 new_stmt.type = left_node.node.type
346
347 attributes.append(
348 util.SQLAlchemyAttribute(
349 name=left_node.name,
350 line=stmt.line,
351 column=stmt.column,
352 typ=left_hand_explicit_type,
353 info=cls.info,
354 )
355 )
356 cls.defs.body[dec_index] = new_stmt
357
358
359def _scan_declarative_assignment_stmt(
360 cls: ClassDef,
361 api: SemanticAnalyzerPluginInterface,
362 stmt: AssignmentStmt,
363 attributes: List[util.SQLAlchemyAttribute],
364) -> None:
365 """Extract mapping information from an assignment statement in a
366 declarative class.
367
368 """
369 lvalue = stmt.lvalues[0]
370 if not isinstance(lvalue, NameExpr):
371 return
372
373 sym = cls.info.names.get(lvalue.name)
374
375 # this establishes that semantic analysis has taken place, which
376 # means the nodes are populated and we are called from an appropriate
377 # hook.
378 assert sym is not None
379 node = sym.node
380
381 if isinstance(node, PlaceholderNode):
382 return
383
384 assert node is lvalue.node
385 assert isinstance(node, Var)
386
387 if node.name == "__abstract__":
388 if api.parse_bool(stmt.rvalue) is True:
389 util.set_is_base(cls.info)
390 return
391 elif node.name == "__tablename__":
392 util.set_has_table(cls.info)
393 elif node.name.startswith("__"):
394 return
395 elif node.name == "_mypy_mapped_attrs":
396 if not isinstance(stmt.rvalue, ListExpr):
397 util.fail(api, "_mypy_mapped_attrs is expected to be a list", stmt)
398 else:
399 for item in stmt.rvalue.items:
400 if isinstance(item, (NameExpr, StrExpr)):
401 apply.apply_mypy_mapped_attr(cls, api, item, attributes)
402
403 left_hand_mapped_type: Optional[Type] = None
404 left_hand_explicit_type: Optional[ProperType] = None
405
406 if node.is_inferred or node.type is None:
407 if isinstance(stmt.type, UnboundType):
408 # look for an explicit Mapped[] type annotation on the left
409 # side with nothing on the right
410
411 # print(stmt.type)
412 # Mapped?[Optional?[A?]]
413
414 left_hand_explicit_type = stmt.type
415
416 if stmt.type.name == "Mapped":
417 mapped_sym = api.lookup_qualified("Mapped", cls)
418 if (
419 mapped_sym is not None
420 and mapped_sym.node is not None
421 and names.type_id_for_named_node(mapped_sym.node)
422 is names.MAPPED
423 ):
424 left_hand_explicit_type = get_proper_type(
425 stmt.type.args[0]
426 )
427 left_hand_mapped_type = stmt.type
428
429 # TODO: do we need to convert from unbound for this case?
430 # left_hand_explicit_type = util._unbound_to_instance(
431 # api, left_hand_explicit_type
432 # )
433 else:
434 node_type = get_proper_type(node.type)
435 if (
436 isinstance(node_type, Instance)
437 and names.type_id_for_named_node(node_type.type) is names.MAPPED
438 ):
439 # print(node.type)
440 # sqlalchemy.orm.attributes.Mapped[<python type>]
441 left_hand_explicit_type = get_proper_type(node_type.args[0])
442 left_hand_mapped_type = node_type
443 else:
444 # print(node.type)
445 # <python type>
446 left_hand_explicit_type = node_type
447 left_hand_mapped_type = None
448
449 if isinstance(stmt.rvalue, TempNode) and left_hand_mapped_type is not None:
450 # annotation without assignment and Mapped is present
451 # as type annotation
452 # equivalent to using _infer_type_from_left_hand_type_only.
453
454 python_type_for_type = left_hand_explicit_type
455 elif isinstance(stmt.rvalue, CallExpr) and isinstance(
456 stmt.rvalue.callee, RefExpr
457 ):
458 python_type_for_type = infer.infer_type_from_right_hand_nameexpr(
459 api, stmt, node, left_hand_explicit_type, stmt.rvalue.callee
460 )
461
462 if python_type_for_type is None:
463 return
464
465 else:
466 return
467
468 assert python_type_for_type is not None
469
470 attributes.append(
471 util.SQLAlchemyAttribute(
472 name=node.name,
473 line=stmt.line,
474 column=stmt.column,
475 typ=python_type_for_type,
476 info=cls.info,
477 )
478 )
479
480 apply.apply_type_to_mapped_statement(
481 api,
482 stmt,
483 lvalue,
484 left_hand_explicit_type,
485 python_type_for_type,
486 )
487
488
489def _scan_for_mapped_bases(
490 cls: ClassDef,
491 api: SemanticAnalyzerPluginInterface,
492) -> None:
493 """Given a class, iterate through its superclass hierarchy to find
494 all other classes that are considered as ORM-significant.
495
496 Locates non-mapped mixins and scans them for mapped attributes to be
497 applied to subclasses.
498
499 """
500
501 info = util.info_for_cls(cls, api)
502
503 if info is None:
504 return
505
506 for base_info in info.mro[1:-1]:
507 if base_info.fullname.startswith("builtins"):
508 continue
509
510 # scan each base for mapped attributes. if they are not already
511 # scanned (but have all their type info), that means they are unmapped
512 # mixins
513 scan_declarative_assignments_and_apply_types(
514 base_info.defn, api, is_mixin_scan=True
515 )
516