Fix type hint for with_for_update() to support tuples of table classes

Fixed typing issue where :meth:`.Select.with_for_update` would not support
lists of ORM entities in the :paramref:`.Select.with_for_update.of`
parameter. Pull request courtesy Shamil.

Fixes: #12730
Closes: #12988
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/12988
Pull-request-sha: 41a38bfe38d8b66da853012e7165cc0cacf7f28a

Change-Id: I61d60a4f4d2b16037da8d5f30e33f5d74fa47374

diff --git a/doc/build/changelog/unreleased_20/12730.rst b/doc/build/changelog/unreleased_20/12730.rst
new file mode 100644 (file)
index 0000000..8d5879d
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/12730.rst
@@ -0,0 +1,8 @@
+.. change::
+    :tags: bug, typing
+    :tickets: 12730
+
+    Fixed typing issue where :meth:`.Select.with_for_update` would not support
+    lists of ORM entities or other FROM clauses in the
+    :paramref:`.Select.with_for_update.of` parameter. Pull request courtesy
+    Shamil.

diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index b8ca7a5659194a3b3b6d596897be3c5992e3fad7..48fc75e3a11b8ea0a7a963cfd93ff67cca381a48 100644 (file)
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -174,13 +174,18 @@ _JoinTargetElement = Union["FromClause", _JoinTargetProtocol]
 _OnClauseElement = Union["ColumnElement[bool]", _JoinTargetProtocol]
 
 _ForUpdateOfArgument = Union[
-    # single column, Table, ORM Entity
+    # single column, Table, ORM entity
     Union[
         "_ColumnExpressionArgument[Any]",
         "_FromClauseArgument",
     ],
-    # or sequence of single column elements
-    Sequence["_ColumnExpressionArgument[Any]"],
+    # or sequence of column, Table, ORM entity
+    Sequence[
+        Union[
+            "_ColumnExpressionArgument[Any]",
+            "_FromClauseArgument",
+        ]
+    ],
 ]
 
 

diff --git a/test/typing/plain_files/sql/common_sql_element.py b/test/typing/plain_files/sql/common_sql_element.py
index e8a10e553db85adb4d238e413774ed068c6add7b..e4aa2cc5a06234e152b854c561668f0f270f8c8b 100644 (file)
--- a/test/typing/plain_files/sql/common_sql_element.py
+++ b/test/typing/plain_files/sql/common_sql_element.py
@@ -52,6 +52,19 @@ class User(Base):
     email: Mapped[str]
 
 
+class A(Base):
+    __tablename__ = "a"
+    id: Mapped[int] = mapped_column(primary_key=True)
+
+
+class B(Base):
+    __tablename__ = "b"
+    id: Mapped[int] = mapped_column(primary_key=True)
+
+
+a_table = Table("a", MetaData(), Column("id", Integer))
+b_table = Table("b", MetaData(), Column("id", Integer))
+
 user_table = Table(
     "user_table", MetaData(), Column("id", Integer), Column("email", String)
 )
@@ -152,6 +165,14 @@ s9174_6 = select(user_table).with_for_update(
     of=[user_table.c.id, user_table.c.email]
 )
 
+# test 12730 - multiple FROM clauses
+s12730_1 = select(A, B).with_for_update(of=(A, B))
+s12730_2 = select(A, B).with_for_update(of=(A,))
+s12730_3 = select(A, B).with_for_update(of=[A, B])
+s12730_4 = select(A, B).with_for_update(of=[A, B])
+s12730_5 = select(a_table, b_table).with_for_update(of=[a_table, b_table])
+
+
 # with_for_update but for query
 session = Session()
 user = session.query(User).with_for_update(of=User)
@@ -162,6 +183,9 @@ user = session.query(user_table).with_for_update(of=user_table.c.id)
 user = session.query(user_table).with_for_update(
     of=[user_table.c.id, user_table.c.email]
 )
+s12730_q1 = session.query(A, B).with_for_update(of=(A, B))
+s12730_q2 = session.query(A, B).with_for_update(of=(A, B))
+
 
 # literal
 assert_type(literal("5"), BindParameter[str])