273 lines
8.8 KiB
Python
273 lines
8.8 KiB
Python
|
# 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 MemberExpr
|
||
|
|
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 get_proper_type
|
||
|
|
from mypy.types import Instance
|
||
|
|
from mypy.types import NoneTyp
|
||
|
|
from mypy.types import ProperType
|
||
|
|
from mypy.types import TypeOfAny
|
||
|
|
from mypy.types import UnboundType
|
||
|
|
from mypy.types import UnionType
|
||
|
|
|
||
|
|
from . import infer
|
||
|
|
from . import util
|
||
|
|
|
||
|
|
|
||
|
|
def _apply_mypy_mapped_attr(
|
||
|
|
cls: ClassDef,
|
||
|
|
api: SemanticAnalyzerPluginInterface,
|
||
|
|
item: Union[NameExpr, StrExpr],
|
||
|
|
cls_metadata: util.DeclClassApplied,
|
||
|
|
) -> None:
|
||
|
|
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 isinstance(stmt.lvalues[0], NameExpr)
|
||
|
|
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 = get_proper_type(stmt.type)
|
||
|
|
assert isinstance(
|
||
|
|
left_hand_explicit_type, (Instance, UnionType, UnboundType)
|
||
|
|
)
|
||
|
|
|
||
|
|
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,
|
||
|
|
) -> None:
|
||
|
|
"""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
|
||
|
|
}
|
||
|
|
update_cls_metadata = False
|
||
|
|
|
||
|
|
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 isinstance(stmt.lvalues[0], NameExpr)
|
||
|
|
and stmt.lvalues[0].name in mapped_attr_lookup
|
||
|
|
and isinstance(stmt.lvalues[0].node, Var)
|
||
|
|
):
|
||
|
|
|
||
|
|
left_node = stmt.lvalues[0].node
|
||
|
|
python_type_for_type = mapped_attr_lookup[stmt.lvalues[0].name]
|
||
|
|
# if we have scanned an UnboundType and now there's a more
|
||
|
|
# specific type than UnboundType, call the re-scan so we
|
||
|
|
# can get that set up correctly
|
||
|
|
if (
|
||
|
|
isinstance(python_type_for_type, UnboundType)
|
||
|
|
and not isinstance(left_node.type, UnboundType)
|
||
|
|
and (
|
||
|
|
isinstance(stmt.rvalue.callee, MemberExpr)
|
||
|
|
and stmt.rvalue.callee.expr.node.fullname
|
||
|
|
== "sqlalchemy.orm.attributes.Mapped"
|
||
|
|
and stmt.rvalue.callee.name == "_empty_constructor"
|
||
|
|
and isinstance(stmt.rvalue.args[0], CallExpr)
|
||
|
|
)
|
||
|
|
):
|
||
|
|
|
||
|
|
python_type_for_type = (
|
||
|
|
infer._infer_type_from_right_hand_nameexpr(
|
||
|
|
api,
|
||
|
|
stmt,
|
||
|
|
left_node,
|
||
|
|
left_node.type,
|
||
|
|
stmt.rvalue.args[0].callee,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
if python_type_for_type is None or isinstance(
|
||
|
|
python_type_for_type, UnboundType
|
||
|
|
):
|
||
|
|
continue
|
||
|
|
|
||
|
|
# update the DeclClassApplied with the better information
|
||
|
|
mapped_attr_lookup[stmt.lvalues[0].name] = python_type_for_type
|
||
|
|
update_cls_metadata = True
|
||
|
|
|
||
|
|
left_node.type = api.named_type(
|
||
|
|
"__sa_Mapped", [python_type_for_type]
|
||
|
|
)
|
||
|
|
|
||
|
|
if update_cls_metadata:
|
||
|
|
cls_metadata.mapped_attr_names[:] = [
|
||
|
|
(k, v) for k, v in mapped_attr_lookup.items()
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def _apply_type_to_mapped_statement(
|
||
|
|
api: SemanticAnalyzerPluginInterface,
|
||
|
|
stmt: AssignmentStmt,
|
||
|
|
lvalue: NameExpr,
|
||
|
|
left_hand_explicit_type: Optional[ProperType],
|
||
|
|
python_type_for_type: Optional[ProperType],
|
||
|
|
) -> 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
|
||
|
|
assert isinstance(left_node, Var)
|
||
|
|
|
||
|
|
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",
|
||
|
|
[] if python_type_for_type is None else [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.attributes.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,
|
||
|
|
) -> None:
|
||
|
|
sym = api.lookup_fully_qualified_or_none(qualified_name)
|
||
|
|
if sym:
|
||
|
|
assert isinstance(sym.node, TypeInfo)
|
||
|
|
type_: ProperType = 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)
|