fix lambda styles for relationship

Change-Id: I810fafc4cbcf0ac9b1f0de764130c81c56367a16

diff --git a/lib/sqlalchemy/orm/relationships.py b/lib/sqlalchemy/orm/relationships.py
index 669a2047c948c7bdfe167230c6acb3a6438700b0..45b9b9beaaee0c5ecfc16bee1094cfbda2ad2e8f 100644 (file)
--- a/lib/sqlalchemy/orm/relationships.py
+++ b/lib/sqlalchemy/orm/relationships.py
@@ -160,9 +160,16 @@ _ORMOrderByArgument = Union[
     Iterable[Union[str, _ColumnExpressionArgument[Any]]],
 ]
 _ORMBackrefArgument = Union[str, Tuple[str, Dict[str, Any]]]
+
+_ORMColCollectionElement = Union[
+    ColumnClause[Any], _HasClauseElement, roles.DMLColumnRole
+]
 _ORMColCollectionArgument = Union[
     str,
-    Sequence[Union[ColumnClause[Any], _HasClauseElement, roles.DMLColumnRole]],
+    Sequence[_ORMColCollectionElement],
+    Callable[[], Sequence[_ORMColCollectionElement]],
+    Callable[[], _ORMColCollectionElement],
+    _ORMColCollectionElement,
 ]
 
 

diff --git a/test/ext/mypy/plain_files/traditional_relationship.py b/test/ext/mypy/plain_files/traditional_relationship.py
index 0edeb694adb2a23789e6a212014b6e3584be5cf2..02afc7c80122b462d2e528fc0049fcdc45c93ef3 100644 (file)
--- a/test/ext/mypy/plain_files/traditional_relationship.py
+++ b/test/ext/mypy/plain_files/traditional_relationship.py
@@ -58,6 +58,25 @@ class Address(Base):
 
     user_style_five = relationship(User, collection_class=set)
 
+    user_fk_style_one: Mapped[List[User]] = relationship(
+        foreign_keys="Address.user_id"
+    )
+    user_fk_style_two: Mapped[List[User]] = relationship(
+        foreign_keys=lambda: Address.user_id
+    )
+    user_fk_style_three: Mapped[List[User]] = relationship(
+        foreign_keys=[user_id]
+    )
+    user_pj_style_one: Mapped[List[User]] = relationship(
+        primaryjoin=user_id == User.id
+    )
+    user_pj_style_two: Mapped[List[User]] = relationship(
+        primaryjoin=lambda: Address.user_id == User.id
+    )
+    user_pj_style_three: Mapped[List[User]] = relationship(
+        primaryjoin="Address.user_id == User.id"
+    )
+
 
 if typing.TYPE_CHECKING:
     # EXPECTED_RE_TYPE: sqlalchemy.*.InstrumentedAttribute\[builtins.list\*?\[traditional_relationship.Address\]\]