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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
|
# ext/mypy/apply.py
# Copyright (C) 2021 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: http://www.opensource.org/licenses/mit-license.php
from typing import Optional
from typing import Union
from mypy import nodes
from mypy.nodes import ARG_NAMED_OPT
from mypy.nodes import Argument
from mypy.nodes import AssignmentStmt
from mypy.nodes import CallExpr
from mypy.nodes import ClassDef
from mypy.nodes import MDEF
from mypy.nodes import NameExpr
from mypy.nodes import StrExpr
from mypy.nodes import SymbolTableNode
from mypy.nodes import TempNode
from mypy.nodes import TypeInfo
from mypy.nodes import Var
from mypy.plugin import SemanticAnalyzerPluginInterface
from mypy.plugins.common import add_method_to_class
from mypy.types import AnyType
from mypy.types import Instance
from mypy.types import NoneTyp
from mypy.types import TypeOfAny
from mypy.types import UnionType
from . import util
def _apply_mypy_mapped_attr(
cls: ClassDef,
api: SemanticAnalyzerPluginInterface,
item: Union[NameExpr, StrExpr],
cls_metadata: util.DeclClassApplied,
):
if isinstance(item, NameExpr):
name = item.name
elif isinstance(item, StrExpr):
name = item.value
else:
return
for stmt in cls.defs.body:
if isinstance(stmt, AssignmentStmt) and stmt.lvalues[0].name == name:
break
else:
util.fail(api, "Can't find mapped attribute {}".format(name), cls)
return
if stmt.type is None:
util.fail(
api,
"Statement linked from _mypy_mapped_attrs has no "
"typing information",
stmt,
)
return
left_hand_explicit_type = stmt.type
cls_metadata.mapped_attr_names.append((name, left_hand_explicit_type))
_apply_type_to_mapped_statement(
api, stmt, stmt.lvalues[0], left_hand_explicit_type, None
)
def _re_apply_declarative_assignments(
cls: ClassDef,
api: SemanticAnalyzerPluginInterface,
cls_metadata: util.DeclClassApplied,
):
"""For multiple class passes, re-apply our left-hand side types as mypy
seems to reset them in place.
"""
mapped_attr_lookup = {
name: typ for name, typ in cls_metadata.mapped_attr_names
}
for stmt in cls.defs.body:
# for a re-apply, all of our statements are AssignmentStmt;
# @declared_attr calls will have been converted and this
# currently seems to be preserved by mypy (but who knows if this
# will change).
if (
isinstance(stmt, AssignmentStmt)
and stmt.lvalues[0].name in mapped_attr_lookup
):
typ = mapped_attr_lookup[stmt.lvalues[0].name]
left_node = stmt.lvalues[0].node
left_node.type = api.named_type("__sa_Mapped", [typ])
def _apply_type_to_mapped_statement(
api: SemanticAnalyzerPluginInterface,
stmt: AssignmentStmt,
lvalue: NameExpr,
left_hand_explicit_type: Optional[Union[Instance, UnionType]],
python_type_for_type: Union[Instance, UnionType],
) -> None:
"""Apply the Mapped[<type>] annotation and right hand object to a
declarative assignment statement.
This converts a Python declarative class statement such as::
class User(Base):
# ...
attrname = Column(Integer)
To one that describes the final Python behavior to Mypy::
class User(Base):
# ...
attrname : Mapped[Optional[int]] = <meaningless temp node>
"""
left_node = lvalue.node
if left_hand_explicit_type is not None:
left_node.type = api.named_type(
"__sa_Mapped", [left_hand_explicit_type]
)
else:
lvalue.is_inferred_def = False
left_node.type = api.named_type("__sa_Mapped", [python_type_for_type])
# so to have it skip the right side totally, we can do this:
# stmt.rvalue = TempNode(AnyType(TypeOfAny.special_form))
# however, if we instead manufacture a new node that uses the old
# one, then we can still get type checking for the call itself,
# e.g. the Column, relationship() call, etc.
# rewrite the node as:
# <attr> : Mapped[<typ>] =
# _sa_Mapped._empty_constructor(<original CallExpr from rvalue>)
# the original right-hand side is maintained so it gets type checked
# internally
column_descriptor = nodes.NameExpr("__sa_Mapped")
column_descriptor.fullname = "sqlalchemy.orm.Mapped"
mm = nodes.MemberExpr(column_descriptor, "_empty_constructor")
orig_call_expr = stmt.rvalue
stmt.rvalue = CallExpr(mm, [orig_call_expr], [nodes.ARG_POS], ["arg1"])
def _add_additional_orm_attributes(
cls: ClassDef,
api: SemanticAnalyzerPluginInterface,
cls_metadata: util.DeclClassApplied,
) -> None:
"""Apply __init__, __table__ and other attributes to the mapped class."""
info = util._info_for_cls(cls, api)
if "__init__" not in info.names and cls_metadata.is_mapped:
mapped_attr_names = {n: t for n, t in cls_metadata.mapped_attr_names}
for mapped_base in cls_metadata.mapped_mro:
base_cls_metadata = util.DeclClassApplied.deserialize(
mapped_base.type.metadata["_sa_decl_class_applied"], api
)
for n, t in base_cls_metadata.mapped_attr_names:
mapped_attr_names.setdefault(n, t)
arguments = []
for name, typ in mapped_attr_names.items():
if typ is None:
typ = AnyType(TypeOfAny.special_form)
arguments.append(
Argument(
variable=Var(name, typ),
type_annotation=typ,
initializer=TempNode(typ),
kind=ARG_NAMED_OPT,
)
)
add_method_to_class(api, cls, "__init__", arguments, NoneTyp())
if "__table__" not in info.names and cls_metadata.has_table:
_apply_placeholder_attr_to_class(
api, cls, "sqlalchemy.sql.schema.Table", "__table__"
)
if cls_metadata.is_mapped:
_apply_placeholder_attr_to_class(
api, cls, "sqlalchemy.orm.mapper.Mapper", "__mapper__"
)
def _apply_placeholder_attr_to_class(
api: SemanticAnalyzerPluginInterface,
cls: ClassDef,
qualified_name: str,
attrname: str,
):
sym = api.lookup_fully_qualified_or_none(qualified_name)
if sym:
assert isinstance(sym.node, TypeInfo)
type_: Union[Instance, AnyType] = Instance(sym.node, [])
else:
type_ = AnyType(TypeOfAny.special_form)
var = Var(attrname)
var.info = cls.info
var.type = type_
cls.info.names[attrname] = SymbolTableNode(MDEF, var)
|
