typing updates to accept with_polymorphic(), aliases

Fixed typing issues where ORM mapped classes and aliased entities could not
be used as keys in result row mappings or as join targets in select
statements. Patterns such as ``row._mapping[User]``,
``row._mapping[aliased(User)]``, ``row._mapping[with_polymorphic(...)]``
(rejected by both mypy and Pylance), and ``.join(aliased(User))``
(rejected by Pylance) are documented and fully supported at runtime but
were previously rejected by type checkers. The type definitions for
:class:`._KeyType` and :class:`._FromClauseArgument` have been updated to
accept these ORM entity types.

Fixes: #13075
Change-Id: Icc3b1ef832b01fd205b1409b2f6d0f211395d4ad

diff --git a/doc/build/changelog/unreleased_20/13075.rst b/doc/build/changelog/unreleased_20/13075.rst
new file mode 100644 (file)
index 0000000..387dde4
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/13075.rst
@@ -0,0 +1,13 @@
+.. change::
+    :tags: bug, typing
+    :tickets: 13075
+
+    Fixed typing issues where ORM mapped classes and aliased entities could not
+    be used as keys in result row mappings or as join targets in select
+    statements. Patterns such as ``row._mapping[User]``,
+    ``row._mapping[aliased(User)]``, ``row._mapping[with_polymorphic(...)]``
+    (rejected by both mypy and Pylance), and ``.join(aliased(User))``
+    (rejected by Pylance) are documented and fully supported at runtime but
+    were previously rejected by type checkers. The type definitions for
+    :class:`._KeyType` and :class:`._FromClauseArgument` have been updated to
+    accept these ORM entity types.

diff --git a/lib/sqlalchemy/engine/result.py b/lib/sqlalchemy/engine/result.py
index 05f7df7c5a8059a7ad958e0c679ac914cf7ef922..a4fd81cf545c1abbc930017acbef6553a43b132c 100644 (file)
--- a/lib/sqlalchemy/engine/result.py
+++ b/lib/sqlalchemy/engine/result.py
@@ -53,10 +53,22 @@ from ..util.typing import TypeVarTuple
 from ..util.typing import Unpack
 
 if typing.TYPE_CHECKING:
+    from typing import Type
+
+    from .. import inspection
+    from ..sql import roles
+    from ..sql._typing import _HasClauseElement
     from ..sql.elements import SQLCoreOperations
     from ..sql.type_api import _ResultProcessorType
 
-_KeyType = Union[str, "SQLCoreOperations[Any]"]
+_KeyType = Union[
+    str,
+    "SQLCoreOperations[Any]",
+    "roles.TypedColumnsClauseRole[Any]",
+    "roles.ColumnsClauseRole",
+    "Type[Any]",
+    "inspection.Inspectable[_HasClauseElement[Any]]",
+]
 _KeyIndexType = Union[_KeyType, int]
 
 # is overridden in cursor using _CursorKeyMapRecType

diff --git a/lib/sqlalchemy/sql/_typing.py b/lib/sqlalchemy/sql/_typing.py
index 71ad7e13eefebee989ca021244b3c52c798b1f20..c4575cdb8c55e0f6a543083f06e38791d894fc3b 100644 (file)
--- a/lib/sqlalchemy/sql/_typing.py
+++ b/lib/sqlalchemy/sql/_typing.py
@@ -242,6 +242,7 @@ _InfoType = Dict[Any, Any]
 
 _FromClauseArgument = Union[
     roles.FromClauseRole,
+    roles.TypedColumnsClauseRole[Any],
     Type[Any],
     Inspectable[_HasClauseElement[Any]],
     _HasClauseElement[Any],

diff --git a/test/typing/plain_files/engine/engine_result.py b/test/typing/plain_files/engine/engine_result.py
index 4c4b030f18c3f763cfaf7b3905e687aa46715367..a70da51a13a703d445013236036565b65dbf3538 100644 (file)
--- a/test/typing/plain_files/engine/engine_result.py
+++ b/test/typing/plain_files/engine/engine_result.py
@@ -26,7 +26,7 @@ def row_one(row: Row[int, str, bool]) -> None:
     assert_type(rm["foo"], Any)
     assert_type(rm[column("bar")], Any)
 
-    # EXPECTED_MYPY_RE: Invalid index type "int" for "RowMapping"; expected type "(str \| SQLCoreOperations\[Any\]|Union\[str, SQLCoreOperations\[Any\]\])"  # noqa: E501
+    # EXPECTED_MYPY_RE: Invalid index type "int" for "RowMapping"; expected type ".*"  # noqa: E501
     rm[3]
 
 

diff --git a/test/typing/plain_files/orm/orm_querying.py b/test/typing/plain_files/orm/orm_querying.py
index 8f18e2fcc18124d18fda3b47e18f4384ba3c2c43..4820ca6e4d6d588883de94abc93caf8640d5056c 100644 (file)
--- a/test/typing/plain_files/orm/orm_querying.py
+++ b/test/typing/plain_files/orm/orm_querying.py
@@ -149,3 +149,26 @@ def test_10937() -> None:
 def test_bundles() -> None:
     b1 = orm.Bundle("b1", A.id, A.data)
     orm.Bundle("b2", A.id, A.data, b1)
+
+
+def test_row_mapping_with_orm_entities() -> None:
+    """Test row._mapping access with ORM entities and aliased entities."""
+    from sqlalchemy import create_engine
+
+    engine = create_engine("sqlite://")
+
+    # Test with regular mapped class
+    stmt1 = select(A, B)
+    with orm.Session(engine) as session:
+        for row in session.execute(stmt1):
+            _ = row._mapping[A]
+            _ = row._mapping[B]
+
+    # Test with aliased class
+    a_alias = aliased(A)
+    b_alias = aliased(B)
+    stmt2 = select(a_alias, b_alias)
+    with orm.Session(engine) as session:
+        for row in session.execute(stmt2):
+            _ = row._mapping[a_alias]
+            _ = row._mapping[b_alias]

diff --git a/test/typing/plain_files/orm/polymorphic.py b/test/typing/plain_files/orm/polymorphic.py
new file mode 100644 (file)
index 0000000..4c532e3
--- /dev/null
+++ b/test/typing/plain_files/orm/polymorphic.py
@@ -0,0 +1,120 @@
+from typing import assert_type
+from typing import TYPE_CHECKING
+
+from sqlalchemy import create_engine
+from sqlalchemy import ForeignKey
+from sqlalchemy import select
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import mapped_column
+from sqlalchemy.orm import Session
+from sqlalchemy.orm import with_polymorphic
+from sqlalchemy.orm.util import AliasedClass
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class Company(Base):
+    __tablename__ = "company"
+    company_id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[str]
+
+
+class Employee(Base):
+    __tablename__ = "employee"
+    __mapper_args__ = {
+        "polymorphic_on": "type",
+        "polymorphic_identity": "employee",
+    }
+    employee_id: Mapped[int] = mapped_column(primary_key=True)
+    company_id: Mapped[int] = mapped_column(ForeignKey("company.company_id"))
+    name: Mapped[str]
+    type: Mapped[str]
+
+
+class Manager(Employee):
+    __tablename__ = "manager"
+    __mapper_args__ = {
+        "polymorphic_identity": "manager",
+    }
+    employee_id: Mapped[int] = mapped_column(
+        ForeignKey("employee.employee_id"), primary_key=True
+    )
+    manager_data: Mapped[str]
+
+
+class Engineer(Employee):
+    __tablename__ = "engineer"
+    __mapper_args__ = {
+        "polymorphic_identity": "engineer",
+    }
+    employee_id: Mapped[int] = mapped_column(
+        ForeignKey("employee.employee_id"), primary_key=True
+    )
+    engineer_info: Mapped[str]
+
+
+engine = create_engine("sqlite://")
+
+
+def test_with_polymorphic_aliased_in_join() -> None:
+    """Test that with_polymorphic with aliased=True can be used in join().
+
+    Relates to discussion #13075.
+    """
+    manager_employee = with_polymorphic(
+        Employee, [Manager], aliased=True, flat=True
+    )
+    engineer_employee = with_polymorphic(
+        Employee, [Engineer], aliased=True, flat=True
+    )
+
+    if TYPE_CHECKING:
+        assert_type(manager_employee, AliasedClass[Employee])
+        assert_type(engineer_employee, AliasedClass[Employee])
+
+    # Should not produce a type error
+    stmt = select(manager_employee, engineer_employee).join(
+        engineer_employee,
+        engineer_employee.company_id == manager_employee.company_id,
+    )
+
+    with Session(engine) as session:
+        session.execute(stmt)
+
+
+def test_with_polymorphic_aliased_in_row_mapping() -> None:
+    """Test that with_polymorphic result can be used as key in row._mapping.
+
+    Relates to discussion #13075.
+    """
+    manager_employee = with_polymorphic(
+        Employee, [Manager], aliased=True, flat=True
+    )
+    engineer_employee = with_polymorphic(
+        Employee, [Engineer], aliased=True, flat=True
+    )
+
+    stmt = select(manager_employee, engineer_employee).join(
+        engineer_employee,
+        engineer_employee.company_id == manager_employee.company_id,
+    )
+
+    with Session(engine) as session:
+        for row in session.execute(stmt):
+            # Should not produce a type error - discussion #13075
+            _ = row._mapping[manager_employee]
+            _ = row._mapping[engineer_employee]
+
+
+def test_with_polymorphic_non_aliased_in_row_mapping() -> None:
+    """Test with_polymorphic without aliased=True in row._mapping."""
+    poly_employee = with_polymorphic(Employee, [Manager, Engineer])
+
+    stmt = select(poly_employee)
+
+    with Session(engine) as session:
+        for row in session.execute(stmt):
+            _ = row._mapping[poly_employee]