From: Mike Bayer Date: Tue, 13 Jan 2026 14:19:14 +0000 (-0500) Subject: typing updates to accept with_polymorphic(), aliases X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=48410f83407661e009326d7170ab79e5163eb8f1;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git 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 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 05f7df7c5a..a4fd81cf54 100644 --- 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 71ad7e13ee..c4575cdb8c 100644 --- 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 4c4b030f18..a70da51a13 100644 --- 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 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]