set type of type_of to be same as input argument

diff --git a/lib/sqlalchemy/orm/attributes.py b/lib/sqlalchemy/orm/attributes.py
index 5b16ce3d6b32bbcd53901a55859098e28f49a481..33cca564927e7d96e3ae001068ef74add5f132b6 100644 (file)
--- 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 5caf57de7bd69b9115fe5c474f6e04bdf910e66c..fe978be7cf307c75259746f1bf14bf697f3558e1 100644 (file)
--- 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()
 
@@ -199,3 +226,5 @@ with Session(e) as s:
     s.execute(select(B).options(joinedload(B.a7)))
     s.execute(select(B).options(joinedload(B.a8)))
     s.execute(select(B).options(joinedload(B.a9)))
+
+