From 94dfe0dfd7f7e69588d06dcf2ca74a3fce6ad4fc Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Fri, 27 Aug 2021 11:39:55 -0400
Subject: [PATCH] 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
---
 doc/build/changelog/unreleased_14/6950.rst |  6 ++++
 lib/sqlalchemy/ext/mypy/infer.py           | 29 +++++++++++++++-----
 lib/sqlalchemy/ext/mypy/names.py           |  5 ++++
 test/ext/mypy/files/t_6950.py              | 32 ++++++++++++++++++++++
 4 files changed, 65 insertions(+), 7 deletions(-)
 create mode 100644 doc/build/changelog/unreleased_14/6950.rst
 create mode 100644 test/ext/mypy/files/t_6950.py

diff --git a/doc/build/changelog/unreleased_14/6950.rst b/doc/build/changelog/unreleased_14/6950.rst
new file mode 100644
index 0000000000..cdfe69e88d
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/6950.rst
@@ -0,0 +1,6 @@
+.. change::
+    :tags: bug, mypy
+    :tickets: 6950
+
+    Fixed issue where mypy plugin would crash when interpreting a
+    ``query_expression()`` construct.
diff --git a/lib/sqlalchemy/ext/mypy/infer.py b/lib/sqlalchemy/ext/mypy/infer.py
index 85a94bba61..52570f772b 100644
--- a/lib/sqlalchemy/ext/mypy/infer.py
+++ b/lib/sqlalchemy/ext/mypy/infer.py
@@ -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]
+
+        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,
+                )
 
-        # look for column_property() / deferred() etc with Column as first
-        # argument
-        if type_id is names.COLUMN:
+    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(
diff --git a/lib/sqlalchemy/ext/mypy/names.py b/lib/sqlalchemy/ext/mypy/names.py
index 22a79e29b9..3dbfcc7703 100644
--- a/lib/sqlalchemy/ext/mypy/names.py
+++ b/lib/sqlalchemy/ext/mypy/names.py
@@ -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"},
+    ),
 }
 
 
diff --git a/test/ext/mypy/files/t_6950.py b/test/ext/mypy/files/t_6950.py
new file mode 100644
index 0000000000..3ebbf66389
--- /dev/null
+++ b/test/ext/mypy/files/t_6950.py
@@ -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)
-- 
2.47.2