AlkantarClanX12
Your IP : 18.188.183.21
# ext/mypy/decl_class.py
# Copyright (C) 2021-2024 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
)