From: Wouter Kayser <wouterkayser@gmail.com>
Date: Sun, 23 Jun 2024 10:18:05 +0000 (-0400)
Subject: set type of type_of to be same as input argument
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=7d8dfa10df3be8d138dd954708efca7d6ed0e503;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

set type of type_of to be same as input argument

Fixes: #11371

Fixes the of_type method so that it does not return a class with unset generic.
See the original issue for a more detailed explanation.

Closes: #11416
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/11416
Pull-request-sha: ed8d0edebb6b84b9dcffcf24c52f113c37e7fedd

Change-Id: I35637491d6d9c573825f6d13299712626dd521c5
---

diff --git a/lib/sqlalchemy/orm/attributes.py b/lib/sqlalchemy/orm/attributes.py
index 5b16ce3d6b..33cca56492 100644
--- a/lib/sqlalchemy/orm/attributes.py
+++ b/lib/sqlalchemy/orm/attributes.py
@@ -401,7 +401,7 @@ class QueryableAttribute(
             parententity=adapt_to_entity,
         )
 
-    def of_type(self, entity: _EntityType[Any]) -> QueryableAttribute[_T]:
+    def of_type(self, entity: _EntityType[_T]) -> QueryableAttribute[_T]:
         return QueryableAttribute(
             self.class_,
             self.key,
diff --git a/test/typing/plain_files/orm/relationship.py b/test/typing/plain_files/orm/relationship.py
index 5caf57de7b..683e347f19 100644
--- a/test/typing/plain_files/orm/relationship.py
+++ b/test/typing/plain_files/orm/relationship.py
@@ -106,6 +106,30 @@ class SelfReferential(Base):
     )
 
 
+class Employee(Base):
+    __tablename__ = "employee"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    team_id: Mapped[int] = mapped_column(ForeignKey("team.id"))
+    team: Mapped["Team"] = relationship(back_populates="employees")
+
+    __mapper_args__ = {
+        "polymorphic_on": "type",
+        "polymorphic_identity": "employee",
+    }
+
+
+class Team(Base):
+    __tablename__ = "team"
+    id: Mapped[int] = mapped_column(primary_key=True)
+    employees: Mapped[list[Employee]] = relationship("Employee")
+
+
+class Engineer(Employee):
+    engineer_info: Mapped[str]
+
+    __mapper_args__ = {"polymorphic_identity": "engineer"}
+
+
 if typing.TYPE_CHECKING:
     # EXPECTED_RE_TYPE: sqlalchemy.*.InstrumentedAttribute\[Union\[builtins.str, None\]\]
     reveal_type(User.extra)
@@ -137,6 +161,9 @@ if typing.TYPE_CHECKING:
     # EXPECTED_RE_TYPE: sqlalchemy.*.InstrumentedAttribute\[builtins.set\*?\[relationship.MoreMail\]\]
     reveal_type(Address.rel_style_one_anno_only)
 
+    # EXPECTED_RE_TYPE: sqlalchemy.*.QueryableAttribute\[relationship.Engineer\]
+    reveal_type(Team.employees.of_type(Engineer))
+
 
 mapper_registry: registry = registry()