From: Wouter Kayser <wouterkayser@gmail.com>
Date: Sun, 26 May 2024 10:49:30 +0000 (+0200)
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=9935607e40b517923c3b3a1fedac3578a89ab458;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

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 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..fe978be7cf 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()
 
@@ -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)))
+
+