Update 2025-04-13_16:25:39
This commit is contained in:
root
324
venv/lib/python3.11/site-packages/sqlalchemy/ext/mypy/apply.py
Normal file
324
venv/lib/python3.11/site-packages/sqlalchemy/ext/mypy/apply.py
Normal file
@ -0,0 +1,324 @@
|
||||
# ext/mypy/apply.py
|
||||
# Copyright (C) 2021-2025 the SQLAlchemy authors and contributors
|
||||
# <see AUTHORS file>
|
||||
#
|
||||
# This module is part of SQLAlchemy and is released under
|
||||
# the MIT License: https://www.opensource.org/licenses/mit-license.php
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List
|
||||
from typing import Optional
|
||||
from typing import Union
|
||||
|
||||
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 RefExpr
|
||||
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
|
||||
from .names import expr_to_mapped_constructor
|
||||
from .names import NAMED_TYPE_SQLA_MAPPED
|
||||
|
||||
|
||||
def apply_mypy_mapped_attr(
|
||||
cls: ClassDef,
|
||||
api: SemanticAnalyzerPluginInterface,
|
||||
item: Union[NameExpr, StrExpr],
|
||||
attributes: List[util.SQLAlchemyAttribute],
|
||||
) -> None:
|
||||
if isinstance(item, NameExpr):
|
||||
name = item.name
|
||||
elif isinstance(item, StrExpr):
|
||||
name = item.value
|
||||
else:
|
||||
return None
|
||||
|
||||
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, f"Can't find mapped attribute {name}", cls)
|
||||
return None
|
||||
|
||||
if stmt.type is None:
|
||||
util.fail(
|
||||
api,
|
||||
"Statement linked from _mypy_mapped_attrs has no "
|
||||
"typing information",
|
||||
stmt,
|
||||
)
|
||||
return None
|
||||
|
||||
left_hand_explicit_type = get_proper_type(stmt.type)
|
||||
assert isinstance(
|
||||
left_hand_explicit_type, (Instance, UnionType, UnboundType)
|
||||
)
|
||||
|
||||
attributes.append(
|
||||
util.SQLAlchemyAttribute(
|
||||
name=name,
|
||||
line=item.line,
|
||||
column=item.column,
|
||||
typ=left_hand_explicit_type,
|
||||
info=cls.info,
|
||||
)
|
||||
)
|
||||
|
||||
apply_type_to_mapped_statement(
|
||||
api, stmt, stmt.lvalues[0], left_hand_explicit_type, None
|
||||
)
|
||||
|
||||
|
||||
def re_apply_declarative_assignments(
|
||||
cls: ClassDef,
|
||||
api: SemanticAnalyzerPluginInterface,
|
||||
attributes: List[util.SQLAlchemyAttribute],
|
||||
) -> None:
|
||||
"""For multiple class passes, re-apply our left-hand side types as mypy
|
||||
seems to reset them in place.
|
||||
|
||||
"""
|
||||
mapped_attr_lookup = {attr.name: attr for attr in attributes}
|
||||
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
|
||||
].type
|
||||
|
||||
left_node_proper_type = get_proper_type(left_node.type)
|
||||
|
||||
# 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_proper_type, UnboundType)
|
||||
and (
|
||||
isinstance(stmt.rvalue, CallExpr)
|
||||
and isinstance(stmt.rvalue.callee, MemberExpr)
|
||||
and isinstance(stmt.rvalue.callee.expr, NameExpr)
|
||||
and stmt.rvalue.callee.expr.node is not None
|
||||
and stmt.rvalue.callee.expr.node.fullname
|
||||
== NAMED_TYPE_SQLA_MAPPED
|
||||
and stmt.rvalue.callee.name == "_empty_constructor"
|
||||
and isinstance(stmt.rvalue.args[0], CallExpr)
|
||||
and isinstance(stmt.rvalue.args[0].callee, RefExpr)
|
||||
)
|
||||
):
|
||||
new_python_type_for_type = (
|
||||
infer.infer_type_from_right_hand_nameexpr(
|
||||
api,
|
||||
stmt,
|
||||
left_node,
|
||||
left_node_proper_type,
|
||||
stmt.rvalue.args[0].callee,
|
||||
)
|
||||
)
|
||||
|
||||
if new_python_type_for_type is not None and not isinstance(
|
||||
new_python_type_for_type, UnboundType
|
||||
):
|
||||
python_type_for_type = new_python_type_for_type
|
||||
|
||||
# update the SQLAlchemyAttribute with the better
|
||||
# information
|
||||
mapped_attr_lookup[stmt.lvalues[0].name].type = (
|
||||
python_type_for_type
|
||||
)
|
||||
|
||||
update_cls_metadata = True
|
||||
|
||||
if (
|
||||
not isinstance(left_node.type, Instance)
|
||||
or left_node.type.type.fullname != NAMED_TYPE_SQLA_MAPPED
|
||||
):
|
||||
assert python_type_for_type is not None
|
||||
left_node.type = api.named_type(
|
||||
NAMED_TYPE_SQLA_MAPPED, [python_type_for_type]
|
||||
)
|
||||
|
||||
if update_cls_metadata:
|
||||
util.set_mapped_attributes(cls.info, attributes)
|
||||
|
||||
|
||||
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::
|
||||
|
||||
... format: off
|
||||
|
||||
class User(Base):
|
||||
# ...
|
||||
|
||||
attrname : Mapped[Optional[int]] = <meaningless temp node>
|
||||
|
||||
... format: on
|
||||
|
||||
"""
|
||||
left_node = lvalue.node
|
||||
assert isinstance(left_node, Var)
|
||||
|
||||
# to be completely honest I have no idea what the difference between
|
||||
# left_node.type and stmt.type is, what it means if these are different
|
||||
# vs. the same, why in order to get tests to pass I have to assign
|
||||
# to stmt.type for the second case and not the first. this is complete
|
||||
# trying every combination until it works stuff.
|
||||
|
||||
if left_hand_explicit_type is not None:
|
||||
lvalue.is_inferred_def = False
|
||||
left_node.type = api.named_type(
|
||||
NAMED_TYPE_SQLA_MAPPED, [left_hand_explicit_type]
|
||||
)
|
||||
else:
|
||||
lvalue.is_inferred_def = False
|
||||
left_node.type = api.named_type(
|
||||
NAMED_TYPE_SQLA_MAPPED,
|
||||
(
|
||||
[AnyType(TypeOfAny.special_form)]
|
||||
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
|
||||
stmt.rvalue = expr_to_mapped_constructor(stmt.rvalue)
|
||||
|
||||
if stmt.type is not None and python_type_for_type is not None:
|
||||
stmt.type = python_type_for_type
|
||||
|
||||
|
||||
def add_additional_orm_attributes(
|
||||
cls: ClassDef,
|
||||
api: SemanticAnalyzerPluginInterface,
|
||||
attributes: List[util.SQLAlchemyAttribute],
|
||||
) -> None:
|
||||
"""Apply __init__, __table__ and other attributes to the mapped class."""
|
||||
|
||||
info = util.info_for_cls(cls, api)
|
||||
|
||||
if info is None:
|
||||
return
|
||||
|
||||
is_base = util.get_is_base(info)
|
||||
|
||||
if "__init__" not in info.names and not is_base:
|
||||
mapped_attr_names = {attr.name: attr.type for attr in attributes}
|
||||
|
||||
for base in info.mro[1:-1]:
|
||||
if "sqlalchemy" not in info.metadata:
|
||||
continue
|
||||
|
||||
base_cls_attributes = util.get_mapped_attributes(base, api)
|
||||
if base_cls_attributes is None:
|
||||
continue
|
||||
|
||||
for attr in base_cls_attributes:
|
||||
mapped_attr_names.setdefault(attr.name, attr.type)
|
||||
|
||||
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 util.get_has_table(info):
|
||||
_apply_placeholder_attr_to_class(
|
||||
api, cls, "sqlalchemy.sql.schema.Table", "__table__"
|
||||
)
|
||||
if not is_base:
|
||||
_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._fullname = cls.fullname + "." + attrname
|
||||
var.info = cls.info
|
||||
var.type = type_
|
||||
cls.info.names[attrname] = SymbolTableNode(MDEF, var)
|
||||
Reference in New Issue
Block a user