516 lines
17 KiB
Python
516 lines
17 KiB
Python
|
# ext/mypy/decl_class.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 AssignmentStmt
|
||
|
|
from mypy.nodes import CallExpr
|
||
|
|
from mypy.nodes import ClassDef
|
||
|
|
from mypy.nodes import Decorator
|
||
|
|
from mypy.nodes import LambdaExpr
|
||
|
|
from mypy.nodes import ListExpr
|
||
|
|
from mypy.nodes import MemberExpr
|
||
|
|
from mypy.nodes import NameExpr
|
||
|
|
from mypy.nodes import PlaceholderNode
|
||
|
|
from mypy.nodes import RefExpr
|
||
|
|
from mypy.nodes import StrExpr
|
||
|
|
from mypy.nodes import SymbolNode
|
||
|
|
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.types import AnyType
|
||
|
|
from mypy.types import CallableType
|
||
|
|
from mypy.types import get_proper_type
|
||
|
|
from mypy.types import Instance
|
||
|
|
from mypy.types import NoneType
|
||
|
|
from mypy.types import ProperType
|
||
|
|
from mypy.types import Type
|
||
|
|
from mypy.types import TypeOfAny
|
||
|
|
from mypy.types import UnboundType
|
||
|
|
from mypy.types import UnionType
|
||
|
|
|
||
|
|
from . import apply
|
||
|
|
from . import infer
|
||
|
|
from . import names
|
||
|
|
from . import util
|
||
|
|
|
||
|
|
|
||
|
|
def scan_declarative_assignments_and_apply_types(
|
||
|
|
cls: ClassDef,
|
||
|
|
api: SemanticAnalyzerPluginInterface,
|
||
|
|
is_mixin_scan: bool = False,
|
||
|
|
) -> Optional[List[util.SQLAlchemyAttribute]]:
|
||
|
|
info = util.info_for_cls(cls, api)
|
||
|
|
|
||
|
|
if info is None:
|
||
|
|
# this can occur during cached passes
|
||
|
|
return None
|
||
|
|
elif cls.fullname.startswith("builtins"):
|
||
|
|
return None
|
||
|
|
|
||
|
|
mapped_attributes: Optional[List[util.SQLAlchemyAttribute]] = (
|
||
|
|
util.get_mapped_attributes(info, api)
|
||
|
|
)
|
||
|
|
|
||
|
|
# used by assign.add_additional_orm_attributes among others
|
||
|
|
util.establish_as_sqlalchemy(info)
|
||
|
|
|
||
|
|
if mapped_attributes is not None:
|
||
|
|
# ensure that a class that's mapped is always picked up by
|
||
|
|
# its mapped() decorator or declarative metaclass before
|
||
|
|
# it would be detected as an unmapped mixin class
|
||
|
|
|
||
|
|
if not is_mixin_scan:
|
||
|
|
# mypy can call us more than once. it then *may* have reset the
|
||
|
|
# left hand side of everything, but not the right that we removed,
|
||
|
|
# removing our ability to re-scan. but we have the types
|
||
|
|
# here, so lets re-apply them, or if we have an UnboundType,
|
||
|
|
# we can re-scan
|
||
|
|
|
||
|
|
apply.re_apply_declarative_assignments(cls, api, mapped_attributes)
|
||
|
|
|
||
|
|
return mapped_attributes
|
||
|
|
|
||
|
|
mapped_attributes = []
|
||
|
|
|
||
|
|
if not cls.defs.body:
|
||
|
|
# when we get a mixin class from another file, the body is
|
||
|
|
# empty (!) but the names are in the symbol table. so use that.
|
||
|
|
|
||
|
|
for sym_name, sym in info.names.items():
|
||
|
|
_scan_symbol_table_entry(
|
||
|
|
cls, api, sym_name, sym, mapped_attributes
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
for stmt in util.flatten_typechecking(cls.defs.body):
|
||
|
|
if isinstance(stmt, AssignmentStmt):
|
||
|
|
_scan_declarative_assignment_stmt(
|
||
|
|
cls, api, stmt, mapped_attributes
|
||
|
|
)
|
||
|
|
elif isinstance(stmt, Decorator):
|
||
|
|
_scan_declarative_decorator_stmt(
|
||
|
|
cls, api, stmt, mapped_attributes
|
||
|
|
)
|
||
|
|
_scan_for_mapped_bases(cls, api)
|
||
|
|
|
||
|
|
if not is_mixin_scan:
|
||
|
|
apply.add_additional_orm_attributes(cls, api, mapped_attributes)
|
||
|
|
|
||
|
|
util.set_mapped_attributes(info, mapped_attributes)
|
||
|
|
|
||
|
|
return mapped_attributes
|
||
|
|
|
||
|
|
|
||
|
|
def _scan_symbol_table_entry(
|
||
|
|
cls: ClassDef,
|
||
|
|
api: SemanticAnalyzerPluginInterface,
|
||
|
|
name: str,
|
||
|
|
value: SymbolTableNode,
|
||
|
|
attributes: List[util.SQLAlchemyAttribute],
|
||
|
|
) -> None:
|
||
|
|
"""Extract mapping information from a SymbolTableNode that's in the
|
||
|
|
type.names dictionary.
|
||
|
|
|
||
|
|
"""
|
||
|
|
value_type = get_proper_type(value.type)
|
||
|
|
if not isinstance(value_type, Instance):
|
||
|
|
return
|
||
|
|
|
||
|
|
left_hand_explicit_type = None
|
||
|
|
type_id = names.type_id_for_named_node(value_type.type)
|
||
|
|
# type_id = names._type_id_for_unbound_type(value.type.type, cls, api)
|
||
|
|
|
||
|
|
err = False
|
||
|
|
|
||
|
|
# TODO: this is nearly the same logic as that of
|
||
|
|
# _scan_declarative_decorator_stmt, likely can be merged
|
||
|
|
if type_id in {
|
||
|
|
names.MAPPED,
|
||
|
|
names.RELATIONSHIP,
|
||
|
|
names.COMPOSITE_PROPERTY,
|
||
|
|
names.MAPPER_PROPERTY,
|
||
|
|
names.SYNONYM_PROPERTY,
|
||
|
|
names.COLUMN_PROPERTY,
|
||
|
|
}:
|
||
|
|
if value_type.args:
|
||
|
|
left_hand_explicit_type = get_proper_type(value_type.args[0])
|
||
|
|
else:
|
||
|
|
err = True
|
||
|
|
elif type_id is names.COLUMN:
|
||
|
|
if not value_type.args:
|
||
|
|
err = True
|
||
|
|
else:
|
||
|
|
typeengine_arg: Union[ProperType, TypeInfo] = get_proper_type(
|
||
|
|
value_type.args[0]
|
||
|
|
)
|
||
|
|
if isinstance(typeengine_arg, Instance):
|
||
|
|
typeengine_arg = typeengine_arg.type
|
||
|
|
|
||
|
|
if isinstance(typeengine_arg, (UnboundType, TypeInfo)):
|
||
|
|
sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
|
||
|
|
if sym is not None and isinstance(sym.node, TypeInfo):
|
||
|
|
if names.has_base_type_id(sym.node, names.TYPEENGINE):
|
||
|
|
left_hand_explicit_type = UnionType(
|
||
|
|
[
|
||
|
|
infer.extract_python_type_from_typeengine(
|
||
|
|
api, sym.node, []
|
||
|
|
),
|
||
|
|
NoneType(),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
util.fail(
|
||
|
|
api,
|
||
|
|
"Column type should be a TypeEngine "
|
||
|
|
"subclass not '{}'".format(sym.node.fullname),
|
||
|
|
value_type,
|
||
|
|
)
|
||
|
|
|
||
|
|
if err:
|
||
|
|
msg = (
|
||
|
|
"Can't infer type from attribute {} on class {}. "
|
||
|
|
"please specify a return type from this function that is "
|
||
|
|
"one of: Mapped[<python type>], relationship[<target class>], "
|
||
|
|
"Column[<TypeEngine>], MapperProperty[<python type>]"
|
||
|
|
)
|
||
|
|
util.fail(api, msg.format(name, cls.name), cls)
|
||
|
|
|
||
|
|
left_hand_explicit_type = AnyType(TypeOfAny.special_form)
|
||
|
|
|
||
|
|
if left_hand_explicit_type is not None:
|
||
|
|
assert value.node is not None
|
||
|
|
attributes.append(
|
||
|
|
util.SQLAlchemyAttribute(
|
||
|
|
name=name,
|
||
|
|
line=value.node.line,
|
||
|
|
column=value.node.column,
|
||
|
|
typ=left_hand_explicit_type,
|
||
|
|
info=cls.info,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _scan_declarative_decorator_stmt(
|
||
|
|
cls: ClassDef,
|
||
|
|
api: SemanticAnalyzerPluginInterface,
|
||
|
|
stmt: Decorator,
|
||
|
|
attributes: List[util.SQLAlchemyAttribute],
|
||
|
|
) -> None:
|
||
|
|
"""Extract mapping information from a @declared_attr in a declarative
|
||
|
|
class.
|
||
|
|
|
||
|
|
E.g.::
|
||
|
|
|
||
|
|
@reg.mapped
|
||
|
|
class MyClass:
|
||
|
|
# ...
|
||
|
|
|
||
|
|
@declared_attr
|
||
|
|
def updated_at(cls) -> Column[DateTime]:
|
||
|
|
return Column(DateTime)
|
||
|
|
|
||
|
|
Will resolve in mypy as::
|
||
|
|
|
||
|
|
@reg.mapped
|
||
|
|
class MyClass:
|
||
|
|
# ...
|
||
|
|
|
||
|
|
updated_at: Mapped[Optional[datetime.datetime]]
|
||
|
|
|
||
|
|
"""
|
||
|
|
for dec in stmt.decorators:
|
||
|
|
if (
|
||
|
|
isinstance(dec, (NameExpr, MemberExpr, SymbolNode))
|
||
|
|
and names.type_id_for_named_node(dec) is names.DECLARED_ATTR
|
||
|
|
):
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
return
|
||
|
|
|
||
|
|
dec_index = cls.defs.body.index(stmt)
|
||
|
|
|
||
|
|
left_hand_explicit_type: Optional[ProperType] = None
|
||
|
|
|
||
|
|
if util.name_is_dunder(stmt.name):
|
||
|
|
# for dunder names like __table_args__, __tablename__,
|
||
|
|
# __mapper_args__ etc., rewrite these as simple assignment
|
||
|
|
# statements; otherwise mypy doesn't like if the decorated
|
||
|
|
# function has an annotation like ``cls: Type[Foo]`` because
|
||
|
|
# it isn't @classmethod
|
||
|
|
any_ = AnyType(TypeOfAny.special_form)
|
||
|
|
left_node = NameExpr(stmt.var.name)
|
||
|
|
left_node.node = stmt.var
|
||
|
|
new_stmt = AssignmentStmt([left_node], TempNode(any_))
|
||
|
|
new_stmt.type = left_node.node.type
|
||
|
|
cls.defs.body[dec_index] = new_stmt
|
||
|
|
return
|
||
|
|
elif isinstance(stmt.func.type, CallableType):
|
||
|
|
func_type = stmt.func.type.ret_type
|
||
|
|
if isinstance(func_type, UnboundType):
|
||
|
|
type_id = names.type_id_for_unbound_type(func_type, cls, api)
|
||
|
|
else:
|
||
|
|
# this does not seem to occur unless the type argument is
|
||
|
|
# incorrect
|
||
|
|
return
|
||
|
|
|
||
|
|
if (
|
||
|
|
type_id
|
||
|
|
in {
|
||
|
|
names.MAPPED,
|
||
|
|
names.RELATIONSHIP,
|
||
|
|
names.COMPOSITE_PROPERTY,
|
||
|
|
names.MAPPER_PROPERTY,
|
||
|
|
names.SYNONYM_PROPERTY,
|
||
|
|
names.COLUMN_PROPERTY,
|
||
|
|
}
|
||
|
|
and func_type.args
|
||
|
|
):
|
||
|
|
left_hand_explicit_type = get_proper_type(func_type.args[0])
|
||
|
|
elif type_id is names.COLUMN and func_type.args:
|
||
|
|
typeengine_arg = func_type.args[0]
|
||
|
|
if isinstance(typeengine_arg, UnboundType):
|
||
|
|
sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
|
||
|
|
if sym is not None and isinstance(sym.node, TypeInfo):
|
||
|
|
if names.has_base_type_id(sym.node, names.TYPEENGINE):
|
||
|
|
left_hand_explicit_type = UnionType(
|
||
|
|
[
|
||
|
|
infer.extract_python_type_from_typeengine(
|
||
|
|
api, sym.node, []
|
||
|
|
),
|
||
|
|
NoneType(),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
util.fail(
|
||
|
|
api,
|
||
|
|
"Column type should be a TypeEngine "
|
||
|
|
"subclass not '{}'".format(sym.node.fullname),
|
||
|
|
func_type,
|
||
|
|
)
|
||
|
|
|
||
|
|
if left_hand_explicit_type is None:
|
||
|
|
# no type on the decorated function. our option here is to
|
||
|
|
# dig into the function body and get the return type, but they
|
||
|
|
# should just have an annotation.
|
||
|
|
msg = (
|
||
|
|
"Can't infer type from @declared_attr on function '{}'; "
|
||
|
|
"please specify a return type from this function that is "
|
||
|
|
"one of: Mapped[<python type>], relationship[<target class>], "
|
||
|
|
"Column[<TypeEngine>], MapperProperty[<python type>]"
|
||
|
|
)
|
||
|
|
util.fail(api, msg.format(stmt.var.name), stmt)
|
||
|
|
|
||
|
|
left_hand_explicit_type = AnyType(TypeOfAny.special_form)
|
||
|
|
|
||
|
|
left_node = NameExpr(stmt.var.name)
|
||
|
|
left_node.node = stmt.var
|
||
|
|
|
||
|
|
# totally feeling around in the dark here as I don't totally understand
|
||
|
|
# the significance of UnboundType. It seems to be something that is
|
||
|
|
# not going to do what's expected when it is applied as the type of
|
||
|
|
# an AssignmentStatement. So do a feeling-around-in-the-dark version
|
||
|
|
# of converting it to the regular Instance/TypeInfo/UnionType structures
|
||
|
|
# we see everywhere else.
|
||
|
|
if isinstance(left_hand_explicit_type, UnboundType):
|
||
|
|
left_hand_explicit_type = get_proper_type(
|
||
|
|
util.unbound_to_instance(api, left_hand_explicit_type)
|
||
|
|
)
|
||
|
|
|
||
|
|
left_node.node.type = api.named_type(
|
||
|
|
names.NAMED_TYPE_SQLA_MAPPED, [left_hand_explicit_type]
|
||
|
|
)
|
||
|
|
|
||
|
|
# this will ignore the rvalue entirely
|
||
|
|
# rvalue = TempNode(AnyType(TypeOfAny.special_form))
|
||
|
|
|
||
|
|
# rewrite the node as:
|
||
|
|
# <attr> : Mapped[<typ>] =
|
||
|
|
# _sa_Mapped._empty_constructor(lambda: <function body>)
|
||
|
|
# the function body is maintained so it gets type checked internally
|
||
|
|
rvalue = names.expr_to_mapped_constructor(
|
||
|
|
LambdaExpr(stmt.func.arguments, stmt.func.body)
|
||
|
|
)
|
||
|
|
|
||
|
|
new_stmt = AssignmentStmt([left_node], rvalue)
|
||
|
|
new_stmt.type = left_node.node.type
|
||
|
|
|
||
|
|
attributes.append(
|
||
|
|
util.SQLAlchemyAttribute(
|
||
|
|
name=left_node.name,
|
||
|
|
line=stmt.line,
|
||
|
|
column=stmt.column,
|
||
|
|
typ=left_hand_explicit_type,
|
||
|
|
info=cls.info,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
cls.defs.body[dec_index] = new_stmt
|
||
|
|
|
||
|
|
|
||
|
|
def _scan_declarative_assignment_stmt(
|
||
|
|
cls: ClassDef,
|
||
|
|
api: SemanticAnalyzerPluginInterface,
|
||
|
|
stmt: AssignmentStmt,
|
||
|
|
attributes: List[util.SQLAlchemyAttribute],
|
||
|
|
) -> None:
|
||
|
|
"""Extract mapping information from an assignment statement in a
|
||
|
|
declarative class.
|
||
|
|
|
||
|
|
"""
|
||
|
|
lvalue = stmt.lvalues[0]
|
||
|
|
if not isinstance(lvalue, NameExpr):
|
||
|
|
return
|
||
|
|
|
||
|
|
sym = cls.info.names.get(lvalue.name)
|
||
|
|
|
||
|
|
# this establishes that semantic analysis has taken place, which
|
||
|
|
# means the nodes are populated and we are called from an appropriate
|
||
|
|
# hook.
|
||
|
|
assert sym is not None
|
||
|
|
node = sym.node
|
||
|
|
|
||
|
|
if isinstance(node, PlaceholderNode):
|
||
|
|
return
|
||
|
|
|
||
|
|
assert node is lvalue.node
|
||
|
|
assert isinstance(node, Var)
|
||
|
|
|
||
|
|
if node.name == "__abstract__":
|
||
|
|
if api.parse_bool(stmt.rvalue) is True:
|
||
|
|
util.set_is_base(cls.info)
|
||
|
|
return
|
||
|
|
elif node.name == "__tablename__":
|
||
|
|
util.set_has_table(cls.info)
|
||
|
|
elif node.name.startswith("__"):
|
||
|
|
return
|
||
|
|
elif node.name == "_mypy_mapped_attrs":
|
||
|
|
if not isinstance(stmt.rvalue, ListExpr):
|
||
|
|
util.fail(api, "_mypy_mapped_attrs is expected to be a list", stmt)
|
||
|
|
else:
|
||
|
|
for item in stmt.rvalue.items:
|
||
|
|
if isinstance(item, (NameExpr, StrExpr)):
|
||
|
|
apply.apply_mypy_mapped_attr(cls, api, item, attributes)
|
||
|
|
|
||
|
|
left_hand_mapped_type: Optional[Type] = None
|
||
|
|
left_hand_explicit_type: Optional[ProperType] = None
|
||
|
|
|
||
|
|
if node.is_inferred or node.type is None:
|
||
|
|
if isinstance(stmt.type, UnboundType):
|
||
|
|
# look for an explicit Mapped[] type annotation on the left
|
||
|
|
# side with nothing on the right
|
||
|
|
|
||
|
|
# print(stmt.type)
|
||
|
|
# Mapped?[Optional?[A?]]
|
||
|
|
|
||
|
|
left_hand_explicit_type = stmt.type
|
||
|
|
|
||
|
|
if stmt.type.name == "Mapped":
|
||
|
|
mapped_sym = api.lookup_qualified("Mapped", cls)
|
||
|
|
if (
|
||
|
|
mapped_sym is not None
|
||
|
|
and mapped_sym.node is not None
|
||
|
|
and names.type_id_for_named_node(mapped_sym.node)
|
||
|
|
is names.MAPPED
|
||
|
|
):
|
||
|
|
left_hand_explicit_type = get_proper_type(
|
||
|
|
stmt.type.args[0]
|
||
|
|
)
|
||
|
|
left_hand_mapped_type = stmt.type
|
||
|
|
|
||
|
|
# TODO: do we need to convert from unbound for this case?
|
||
|
|
# left_hand_explicit_type = util._unbound_to_instance(
|
||
|
|
# api, left_hand_explicit_type
|
||
|
|
# )
|
||
|
|
else:
|
||
|
|
node_type = get_proper_type(node.type)
|
||
|
|
if (
|
||
|
|
isinstance(node_type, Instance)
|
||
|
|
and names.type_id_for_named_node(node_type.type) is names.MAPPED
|
||
|
|
):
|
||
|
|
# print(node.type)
|
||
|
|
# sqlalchemy.orm.attributes.Mapped[<python type>]
|
||
|
|
left_hand_explicit_type = get_proper_type(node_type.args[0])
|
||
|
|
left_hand_mapped_type = node_type
|
||
|
|
else:
|
||
|
|
# print(node.type)
|
||
|
|
# <python type>
|
||
|
|
left_hand_explicit_type = node_type
|
||
|
|
left_hand_mapped_type = None
|
||
|
|
|
||
|
|
if isinstance(stmt.rvalue, TempNode) and left_hand_mapped_type is not None:
|
||
|
|
# annotation without assignment and Mapped is present
|
||
|
|
# as type annotation
|
||
|
|
# equivalent to using _infer_type_from_left_hand_type_only.
|
||
|
|
|
||
|
|
python_type_for_type = left_hand_explicit_type
|
||
|
|
elif isinstance(stmt.rvalue, CallExpr) and isinstance(
|
||
|
|
stmt.rvalue.callee, RefExpr
|
||
|
|
):
|
||
|
|
python_type_for_type = infer.infer_type_from_right_hand_nameexpr(
|
||
|
|
api, stmt, node, left_hand_explicit_type, stmt.rvalue.callee
|
||
|
|
)
|
||
|
|
|
||
|
|
if python_type_for_type is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
else:
|
||
|
|
return
|
||
|
|
|
||
|
|
assert python_type_for_type is not None
|
||
|
|
|
||
|
|
attributes.append(
|
||
|
|
util.SQLAlchemyAttribute(
|
||
|
|
name=node.name,
|
||
|
|
line=stmt.line,
|
||
|
|
column=stmt.column,
|
||
|
|
typ=python_type_for_type,
|
||
|
|
info=cls.info,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
apply.apply_type_to_mapped_statement(
|
||
|
|
api,
|
||
|
|
stmt,
|
||
|
|
lvalue,
|
||
|
|
left_hand_explicit_type,
|
||
|
|
python_type_for_type,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _scan_for_mapped_bases(
|
||
|
|
cls: ClassDef,
|
||
|
|
api: SemanticAnalyzerPluginInterface,
|
||
|
|
) -> None:
|
||
|
|
"""Given a class, iterate through its superclass hierarchy to find
|
||
|
|
all other classes that are considered as ORM-significant.
|
||
|
|
|
||
|
|
Locates non-mapped mixins and scans them for mapped attributes to be
|
||
|
|
applied to subclasses.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
info = util.info_for_cls(cls, api)
|
||
|
|
|
||
|
|
if info is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
for base_info in info.mro[1:-1]:
|
||
|
|
if base_info.fullname.startswith("builtins"):
|
||
|
|
continue
|
||
|
|
|
||
|
|
# scan each base for mapped attributes. if they are not already
|
||
|
|
# scanned (but have all their type info), that means they are unmapped
|
||
|
|
# mixins
|
||
|
|
scan_declarative_assignments_and_apply_types(
|
||
|
|
base_info.defn, api, is_mixin_scan=True
|
||
|
|
)
|