dont assume argument lists for column property

Fixed issue where mypy plugin would crash when interpreting a
``query_expression()`` construct.

Fixes: #6950
Change-Id: Ic1f28d135bf6eb05c92061430c0d5a3663b804ef
This commit is contained in:
Mike Bayer
2021-08-27 11:39:55 -04:00
parent e2d9ef3fe6
commit 94dfe0dfd7
4 changed files with 65 additions and 7 deletions
+6
View File
@@ -0,0 +1,6 @@
.. change::
:tags: bug, mypy
:tickets: 6950
Fixed issue where mypy plugin would crash when interpreting a
``query_expression()`` construct.
+22 -7
View File
@@ -284,20 +284,35 @@ def _infer_type_from_decl_column_property(
"""
assert isinstance(stmt.rvalue, CallExpr)
first_prop_arg = stmt.rvalue.args[0]
if isinstance(first_prop_arg, CallExpr):
type_id = names.type_id_for_callee(first_prop_arg.callee)
if stmt.rvalue.args:
first_prop_arg = stmt.rvalue.args[0]
# look for column_property() / deferred() etc with Column as first
# argument
if type_id is names.COLUMN:
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,
right_hand_expression=first_prop_arg,
)
return infer_type_from_left_hand_type_only(
+5
View File
@@ -45,6 +45,7 @@ MAPPER_PROPERTY: int = util.symbol("MAPPER_PROPERTY") # type: ignore
AS_DECLARATIVE: int = util.symbol("AS_DECLARATIVE") # type: ignore
AS_DECLARATIVE_BASE: int = util.symbol("AS_DECLARATIVE_BASE") # type: ignore
DECLARATIVE_MIXIN: int = util.symbol("DECLARATIVE_MIXIN") # type: ignore
QUERY_EXPRESSION: int = util.symbol("QUERY_EXPRESSION") # type: ignore
_lookup: Dict[str, Tuple[int, Set[str]]] = {
"Column": (
@@ -150,6 +151,10 @@ _lookup: Dict[str, Tuple[int, Set[str]]] = {
"sqlalchemy.orm.declarative_mixin",
},
),
"query_expression": (
QUERY_EXPRESSION,
{"sqlalchemy.orm.query_expression"},
),
}
+32
View File
@@ -0,0 +1,32 @@
from typing import cast
from sqlalchemy import Column
from sqlalchemy import Integer
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import query_expression
from sqlalchemy.orm import Session
from sqlalchemy.orm import with_expression
Base = declarative_base()
class User(Base):
__tablename__ = "users"
id = Column(Integer, primary_key=True)
foo = Column(Integer)
question_count: Mapped[int] = query_expression()
answer_count: int = query_expression()
s = Session()
q = s.query(User).options(with_expression(User.question_count, User.foo + 5))
u1: User = cast(User, q.first())
qc: int = u1.question_count
print(qc)