AlkantarClanX12
Your IP : 3.128.171.192
# ext/mypy/infer.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 Optional
from typing import Sequence
from mypy.maptype import map_instance_to_supertype
from mypy.nodes import AssignmentStmt
from mypy.nodes import CallExpr
from mypy.nodes import Expression
from mypy.nodes import FuncDef
from mypy.nodes import LambdaExpr
from mypy.nodes import MemberExpr
from mypy.nodes import NameExpr
from mypy.nodes import RefExpr
from mypy.nodes import StrExpr
from mypy.nodes import TypeInfo
from mypy.nodes import Var
from mypy.plugin import SemanticAnalyzerPluginInterface
from mypy.subtypes import is_subtype
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 TypeOfAny
from mypy.types import UnionType
from . import names
from . import util
def infer_type_from_right_hand_nameexpr(
api: SemanticAnalyzerPluginInterface,
stmt: AssignmentStmt,
node: Var,
left_hand_explicit_type: Optional[ProperType],
infer_from_right_side: RefExpr,
) -> Optional[ProperType]:
type_id = names.type_id_for_callee(infer_from_right_side)
if type_id is None:
return None
elif type_id is names.MAPPED:
python_type_for_type = _infer_type_from_mapped(
api, stmt, node, left_hand_explicit_type, infer_from_right_side
)
elif type_id is names.COLUMN:
python_type_for_type = _infer_type_from_decl_column(
api, stmt, node, left_hand_explicit_type
)
elif type_id is names.RELATIONSHIP:
python_type_for_type = _infer_type_from_relationship(
api, stmt, node, left_hand_explicit_type
)
elif type_id is names.COLUMN_PROPERTY:
python_type_for_type = _infer_type_from_decl_column_property(
api, stmt, node, left_hand_explicit_type
)
elif type_id is names.SYNONYM_PROPERTY:
python_type_for_type = infer_type_from_left_hand_type_only(
api, node, left_hand_explicit_type
)
elif type_id is names.COMPOSITE_PROPERTY:
python_type_for_type = _infer_type_from_decl_composite_property(
api, stmt, node, left_hand_explicit_type
)
else:
return None
return python_type_for_type
def _infer_type_from_relationship(
api: SemanticAnalyzerPluginInterface,
stmt: AssignmentStmt,
node: Var,
left_hand_explicit_type: Optional[ProperType],
) -> Optional[ProperType]:
"""Infer the type of mapping from a relationship.
E.g.::
@reg.mapped
class MyClass:
# ...
addresses = relationship(Address, uselist=True)
order: Mapped["Order"] = relationship("Order")
Will resolve in mypy as::
@reg.mapped
class MyClass:
# ...
addresses: Mapped[List[Address]]
order: Mapped["Order"]
"""
assert isinstance(stmt.rvalue, CallExpr)
target_cls_arg = stmt.rvalue.args[0]
python_type_for_type: Optional[ProperType] = None
if isinstance(target_cls_arg, NameExpr) and isinstance(
target_cls_arg.node, TypeInfo
):
# type
related_object_type = target_cls_arg.node
python_type_for_type = Instance(related_object_type, [])
# other cases not covered - an error message directs the user
# to set an explicit type annotation
#
# node.type == str, it's a string
# if isinstance(target_cls_arg, NameExpr) and isinstance(
# target_cls_arg.node, Var
# )
# points to a type
# isinstance(target_cls_arg, NameExpr) and isinstance(
# target_cls_arg.node, TypeAlias
# )
# string expression
# isinstance(target_cls_arg, StrExpr)
uselist_arg = util.get_callexpr_kwarg(stmt.rvalue, "uselist")
collection_cls_arg: Optional[Expression] = util.get_callexpr_kwarg(
stmt.rvalue, "collection_class"
)
type_is_a_collection = False
# this can be used to determine Optional for a many-to-one
# in the same way nullable=False could be used, if we start supporting
# that.
# innerjoin_arg = util.get_callexpr_kwarg(stmt.rvalue, "innerjoin")
if (
uselist_arg is not None
and api.parse_bool(uselist_arg) is True
and collection_cls_arg is None
):
type_is_a_collection = True
if python_type_for_type is not None:
python_type_for_type = api.named_type(
names.NAMED_TYPE_BUILTINS_LIST, [python_type_for_type]
)
elif (
uselist_arg is None or api.parse_bool(uselist_arg) is True
) and collection_cls_arg is not None:
type_is_a_collection = True
if isinstance(collection_cls_arg, CallExpr):
collection_cls_arg = collection_cls_arg.callee
if isinstance(collection_cls_arg, NameExpr) and isinstance(
collection_cls_arg.node, TypeInfo
):
if python_type_for_type is not None:
# this can still be overridden by the left hand side
# within _infer_Type_from_left_and_inferred_right
python_type_for_type = Instance(
collection_cls_arg.node, [python_type_for_type]
)
elif (
isinstance(collection_cls_arg, NameExpr)
and isinstance(collection_cls_arg.node, FuncDef)
and collection_cls_arg.node.type is not None
):
if python_type_for_type is not None:
# this can still be overridden by the left hand side
# within _infer_Type_from_left_and_inferred_right
# TODO: handle mypy.types.Overloaded
if isinstance(collection_cls_arg.node.type, CallableType):
rt = get_proper_type(collection_cls_arg.node.type.ret_type)
if isinstance(rt, CallableType):
callable_ret_type = get_proper_type(rt.ret_type)
if isinstance(callable_ret_type, Instance):
python_type_for_type = Instance(
callable_ret_type.type,
[python_type_for_type],
)
else:
util.fail(
api,
"Expected Python collection type for "
"collection_class parameter",
stmt.rvalue,
)
python_type_for_type = None
elif uselist_arg is not None and api.parse_bool(uselist_arg) is False:
if collection_cls_arg is not None:
util.fail(
api,
"Sending uselist=False and collection_class at the same time "
"does not make sense",
stmt.rvalue,
)
if python_type_for_type is not None:
python_type_for_type = UnionType(
[python_type_for_type, NoneType()]
)
else:
if left_hand_explicit_type is None:
msg = (
"Can't infer scalar or collection for ORM mapped expression "
"assigned to attribute '{}' if both 'uselist' and "
"'collection_class' arguments are absent from the "
"relationship(); please specify a "
"type annotation on the left hand side."
)
util.fail(api, msg.format(node.name), node)
if python_type_for_type is None:
return infer_type_from_left_hand_type_only(
api, node, left_hand_explicit_type
)
elif left_hand_explicit_type is not None:
if type_is_a_collection:
assert isinstance(left_hand_explicit_type, Instance)
assert isinstance(python_type_for_type, Instance)
return _infer_collection_type_from_left_and_inferred_right(
api, node, left_hand_explicit_type, python_type_for_type
)
else:
return _infer_type_from_left_and_inferred_right(
api,
node,
left_hand_explicit_type,
python_type_for_type,
)
else:
return python_type_for_type
def _infer_type_from_decl_composite_property(
api: SemanticAnalyzerPluginInterface,
stmt: AssignmentStmt,
node: Var,
left_hand_explicit_type: Optional[ProperType],
) -> Optional[ProperType]:
"""Infer the type of mapping from a Composite."""
assert isinstance(stmt.rvalue, CallExpr)
target_cls_arg = stmt.rvalue.args[0]
python_type_for_type = None
if isinstance(target_cls_arg, NameExpr) and isinstance(
target_cls_arg.node, TypeInfo
):
related_object_type = target_cls_arg.node
python_type_for_type = Instance(related_object_type, [])
else:
python_type_for_type = None
if python_type_for_type is None:
return infer_type_from_left_hand_type_only(
api, node, left_hand_explicit_type
)
elif left_hand_explicit_type is not None:
return _infer_type_from_left_and_inferred_right(
api, node, left_hand_explicit_type, python_type_for_type
)
else:
return python_type_for_type
def _infer_type_from_mapped(
api: SemanticAnalyzerPluginInterface,
stmt: AssignmentStmt,
node: Var,
left_hand_explicit_type: Optional[ProperType],
infer_from_right_side: RefExpr,
) -> Optional[ProperType]:
"""Infer the type of mapping from a right side expression
that returns Mapped.
"""
assert isinstance(stmt.rvalue, CallExpr)
# (Pdb) print(stmt.rvalue.callee)
# NameExpr(query_expression [sqlalchemy.orm._orm_constructors.query_expression]) # noqa: E501
# (Pdb) stmt.rvalue.callee.node
# <mypy.nodes.FuncDef object at 0x7f8d92fb5940>
# (Pdb) stmt.rvalue.callee.node.type
# def [_T] (default_expr: sqlalchemy.sql.elements.ColumnElement[_T`-1] =) -> sqlalchemy.orm.base.Mapped[_T`-1] # noqa: E501
# sqlalchemy.orm.base.Mapped[_T`-1]
# the_mapped_type = stmt.rvalue.callee.node.type.ret_type
# TODO: look at generic ref and either use that,
# or reconcile w/ what's present, etc.
the_mapped_type = util.type_for_callee(infer_from_right_side) # noqa
return infer_type_from_left_hand_type_only(
api, node, left_hand_explicit_type
)
def _infer_type_from_decl_column_property(
api: SemanticAnalyzerPluginInterface,
stmt: AssignmentStmt,
node: Var,
left_hand_explicit_type: Optional[ProperType],
) -> Optional[ProperType]:
"""Infer the type of mapping from a ColumnProperty.
This includes mappings against ``column_property()`` as well as the
``deferred()`` function.
"""
assert isinstance(stmt.rvalue, CallExpr)
if stmt.rvalue.args:
first_prop_arg = stmt.rvalue.args[0]
if isinstance(first_prop_arg, CallExpr):
type_id = names.type_id_for_callee(first_prop_arg.callee)
# look for column_property() / deferred() etc with Column as first
# argument
if type_id is names.COLUMN:
return _infer_type_from_decl_column(
api,
stmt,
node,
left_hand_explicit_type,
right_hand_expression=first_prop_arg,
)
if isinstance(stmt.rvalue, CallExpr):
type_id = names.type_id_for_callee(stmt.rvalue.callee)
# this is probably not strictly necessary as we have to use the left
# hand type for query expression in any case. any other no-arg
# column prop objects would go here also
if type_id is names.QUERY_EXPRESSION:
return _infer_type_from_decl_column(
api,
stmt,
node,
left_hand_explicit_type,
)
return infer_type_from_left_hand_type_only(
api, node, left_hand_explicit_type
)
def _infer_type_from_decl_column(
api: SemanticAnalyzerPluginInterface,
stmt: AssignmentStmt,
node: Var,
left_hand_explicit_type: Optional[ProperType],
right_hand_expression: Optional[CallExpr] = None,
) -> Optional[ProperType]:
"""Infer the type of mapping from a Column.
E.g.::
@reg.mapped
class MyClass:
# ...
a = Column(Integer)
b = Column("b", String)
c: Mapped[int] = Column(Integer)
d: bool = Column(Boolean)
Will resolve in MyPy as::
@reg.mapped
class MyClass:
# ...
a : Mapped[int]
b : Mapped[str]
c: Mapped[int]
d: Mapped[bool]
"""
assert isinstance(node, Var)
callee = None
if right_hand_expression is None:
if not isinstance(stmt.rvalue, CallExpr):
return None
right_hand_expression = stmt.rvalue
for column_arg in right_hand_expression.args[0:2]:
if isinstance(column_arg, CallExpr):
if isinstance(column_arg.callee, RefExpr):
# x = Column(String(50))
callee = column_arg.callee
type_args: Sequence[Expression] = column_arg.args
break
elif isinstance(column_arg, (NameExpr, MemberExpr)):
if isinstance(column_arg.node, TypeInfo):
# x = Column(String)
callee = column_arg
type_args = ()
break
else:
# x = Column(some_name, String), go to next argument
continue
elif isinstance(column_arg, (StrExpr,)):
# x = Column("name", String), go to next argument
continue
elif isinstance(column_arg, (LambdaExpr,)):
# x = Column("name", String, default=lambda: uuid.uuid4())
# go to next argument
continue
else:
assert False
if callee is None:
return None
if isinstance(callee.node, TypeInfo) and names.mro_has_id(
callee.node.mro, names.TYPEENGINE
):
python_type_for_type = extract_python_type_from_typeengine(
api, callee.node, type_args
)
if left_hand_explicit_type is not None:
return _infer_type_from_left_and_inferred_right(
api, node, left_hand_explicit_type, python_type_for_type
)
else:
return UnionType([python_type_for_type, NoneType()])
else:
# it's not TypeEngine, it's typically implicitly typed
# like ForeignKey. we can't infer from the right side.
return infer_type_from_left_hand_type_only(
api, node, left_hand_explicit_type
)
def _infer_type_from_left_and_inferred_right(
api: SemanticAnalyzerPluginInterface,
node: Var,
left_hand_explicit_type: ProperType,
python_type_for_type: ProperType,
orig_left_hand_type: Optional[ProperType] = None,
orig_python_type_for_type: Optional[ProperType] = None,
) -> Optional[ProperType]:
"""Validate type when a left hand annotation is present and we also
could infer the right hand side::
attrname: SomeType = Column(SomeDBType)
"""
if orig_left_hand_type is None:
orig_left_hand_type = left_hand_explicit_type
if orig_python_type_for_type is None:
orig_python_type_for_type = python_type_for_type
if not is_subtype(left_hand_explicit_type, python_type_for_type):
effective_type = api.named_type(
names.NAMED_TYPE_SQLA_MAPPED, [orig_python_type_for_type]
)
msg = (
"Left hand assignment '{}: {}' not compatible "
"with ORM mapped expression of type {}"
)
util.fail(
api,
msg.format(
node.name,
util.format_type(orig_left_hand_type, api.options),
util.format_type(effective_type, api.options),
),
node,
)
return orig_left_hand_type
def _infer_collection_type_from_left_and_inferred_right(
api: SemanticAnalyzerPluginInterface,
node: Var,
left_hand_explicit_type: Instance,
python_type_for_type: Instance,
) -> Optional[ProperType]:
orig_left_hand_type = left_hand_explicit_type
orig_python_type_for_type = python_type_for_type
if left_hand_explicit_type.args:
left_hand_arg = get_proper_type(left_hand_explicit_type.args[0])
python_type_arg = get_proper_type(python_type_for_type.args[0])
else:
left_hand_arg = left_hand_explicit_type
python_type_arg = python_type_for_type
assert isinstance(left_hand_arg, (Instance, UnionType))
assert isinstance(python_type_arg, (Instance, UnionType))
return _infer_type_from_left_and_inferred_right(
api,
node,
left_hand_arg,
python_type_arg,
orig_left_hand_type=orig_left_hand_type,
orig_python_type_for_type=orig_python_type_for_type,
)
def infer_type_from_left_hand_type_only(
api: SemanticAnalyzerPluginInterface,
node: Var,
left_hand_explicit_type: Optional[ProperType],
) -> Optional[ProperType]:
"""Determine the type based on explicit annotation only.
if no annotation were present, note that we need one there to know
the type.
"""
if left_hand_explicit_type is None:
msg = (
"Can't infer type from ORM mapped expression "
"assigned to attribute '{}'; please specify a "
"Python type or "
"Mapped[<python type>] on the left hand side."
)
util.fail(api, msg.format(node.name), node)
return api.named_type(
names.NAMED_TYPE_SQLA_MAPPED, [AnyType(TypeOfAny.special_form)]
)
else:
# use type from the left hand side
return left_hand_explicit_type
def extract_python_type_from_typeengine(
api: SemanticAnalyzerPluginInterface,
node: TypeInfo,
type_args: Sequence[Expression],
) -> ProperType:
if node.fullname == "sqlalchemy.sql.sqltypes.Enum" and type_args:
first_arg = type_args[0]
if isinstance(first_arg, RefExpr) and isinstance(
first_arg.node, TypeInfo
):
for base_ in first_arg.node.mro:
if base_.fullname == "enum.Enum":
return Instance(first_arg.node, [])
# TODO: support other pep-435 types here
else:
return api.named_type(names.NAMED_TYPE_BUILTINS_STR, [])
assert node.has_base("sqlalchemy.sql.type_api.TypeEngine"), (
"could not extract Python type from node: %s" % node
)
type_engine_sym = api.lookup_fully_qualified_or_none(
"sqlalchemy.sql.type_api.TypeEngine"
)
assert type_engine_sym is not None and isinstance(
type_engine_sym.node, TypeInfo
)
type_engine = map_instance_to_supertype(
Instance(node, []),
type_engine_sym.node,
)
return get_proper_type(type_engine.args[-1])