From 64b3068178b8406caee011d58ecfef75626ea5ea Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Tue, 13 Jan 2026 09:19:14 -0500 Subject: [PATCH] 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 (cherry picked from commit 48410f83407661e009326d7170ab79e5163eb8f1) --- doc/build/changelog/unreleased_20/13075.rst | 13 ++ lib/sqlalchemy/engine/result.py | 14 +- lib/sqlalchemy/sql/_typing.py | 1 + .../plain_files/engine/engine_result.py | 2 +- test/typing/plain_files/orm/orm_querying.py | 23 ++++ test/typing/plain_files/orm/polymorphic.py | 120 ++++++++++++++++++ 6 files changed, 171 insertions(+), 2 deletions(-) create mode 100644 doc/build/changelog/unreleased_20/13075.rst create mode 100644 test/typing/plain_files/orm/polymorphic.py diff --git a/doc/build/changelog/unreleased_20/13075.rst b/doc/build/changelog/unreleased_20/13075.rst new file mode 100644 index 0000000000..387dde4b39 --- /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 c54a965c08..dc7b0727c6 100644 --- a/lib/sqlalchemy/engine/result.py +++ b/lib/sqlalchemy/engine/result.py @@ -52,10 +52,22 @@ else: from sqlalchemy.cyextension.resultproxy import tuplegetter as tuplegetter 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 18639df2f5..7df3a0c63f 100644 --- a/lib/sqlalchemy/sql/_typing.py +++ b/lib/sqlalchemy/sql/_typing.py @@ -241,6 +241,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 8216a69bb5..72a4761ad7 100644 --- a/test/typing/plain_files/engine/engine_result.py +++ b/test/typing/plain_files/engine/engine_result.py @@ -24,7 +24,7 @@ def row_one(row: Row[Tuple[int, str, bool]]) -> None: # EXPECTED_TYPE: Any reveal_type(rm[column("bar")]) - # 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 8f18e2fcc1..4820ca6e4d 100644 --- 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 index 0000000000..4c532e3346 --- /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] -- 2.47.3