From: Yurii Karabas <1998uriyyo@gmail.com>
Date: Sat, 16 Dec 2023 20:38:05 +0000 (+0200)
Subject: Add extra arg for tuple overloads
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=11f5698093c1e3e2763b3b0c0d5c3b05dc19ca32;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

Add extra arg for tuple overloads
---

diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index d18f9ad5cd..6da6f2af1a 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -1575,7 +1575,10 @@ class Query(
         __ent6: _TCCA[_T6],
         __ent7: _TCCA[_T7],
         /,
-    ) -> RowReturningQuery[_T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7]:
+        *entities: _ColumnsClauseArgument[Any],
+    ) -> RowReturningQuery[
+        _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, Unpack[TupleAny]
+    ]:
         ...
 
     # END OVERLOADED FUNCTIONS self.with_entities
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index 092bde7ecc..eb45fb9dbb 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -2892,7 +2892,10 @@ class Session(_SessionClassMethods, EventTarget):
         __ent6: _TCCA[_T6],
         __ent7: _TCCA[_T7],
         /,
-    ) -> RowReturningQuery[_T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7]:
+        *entities: _ColumnsClauseArgument[Any],
+    ) -> RowReturningQuery[
+        _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, Unpack[TupleAny]
+    ]:
         ...
 
     # END OVERLOADED FUNCTIONS self.query
diff --git a/lib/sqlalchemy/sql/_selectable_constructors.py b/lib/sqlalchemy/sql/_selectable_constructors.py
index 2df2ef1cb6..c8a31b643e 100644
--- a/lib/sqlalchemy/sql/_selectable_constructors.py
+++ b/lib/sqlalchemy/sql/_selectable_constructors.py
@@ -442,7 +442,10 @@ def select(
     __ent8: _TCCA[_T8],
     __ent9: _TCCA[_T9],
     /,
-) -> Select[_T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, _T8, _T9]:
+    *entities: _ColumnsClauseArgument[Any],
+) -> Select[
+    _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, _T8, _T9, Unpack[TupleAny]
+]:
     ...
 
 
diff --git a/lib/sqlalchemy/sql/dml.py b/lib/sqlalchemy/sql/dml.py
index cf605ae4c7..8d9f995c4c 100644
--- a/lib/sqlalchemy/sql/dml.py
+++ b/lib/sqlalchemy/sql/dml.py
@@ -1399,9 +1399,11 @@ class Insert(ValuesBase):
             __ent6: _TCCA[_T6],
             __ent7: _TCCA[_T7],
             /,
-            *,
+            *entities: _ColumnsClauseArgument[Any],
             sort_by_parameter_order: bool = False,
-        ) -> ReturningInsert[_T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7]:
+        ) -> ReturningInsert[
+            _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, Unpack[TupleAny]
+        ]:
             ...
 
         # END OVERLOADED FUNCTIONS self.returning
@@ -1688,7 +1690,10 @@ class Update(DMLWhereBase, ValuesBase):
             __ent6: _TCCA[_T6],
             __ent7: _TCCA[_T7],
             /,
-        ) -> ReturningUpdate[_T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7]:
+            *entities: _ColumnsClauseArgument[Any],
+        ) -> ReturningUpdate[
+            _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, Unpack[TupleAny]
+        ]:
             ...
 
         # END OVERLOADED FUNCTIONS self.returning
@@ -1831,7 +1836,10 @@ class Delete(DMLWhereBase, UpdateBase):
             __ent6: _TCCA[_T6],
             __ent7: _TCCA[_T7],
             /,
-        ) -> ReturningDelete[_T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7]:
+            *entities: _ColumnsClauseArgument[Any],
+        ) -> ReturningDelete[
+            _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, Unpack[TupleAny]
+        ]:
             ...
 
         # END OVERLOADED FUNCTIONS self.returning
diff --git a/tools/generate_tuple_map_overloads.py b/tools/generate_tuple_map_overloads.py
index 0a300c20ee..9ca648333c 100644
--- a/tools/generate_tuple_map_overloads.py
+++ b/tools/generate_tuple_map_overloads.py
@@ -82,17 +82,26 @@ def process_module(modname: str, filename: str, cmd: code_writer_cmd) -> str:
                 )
 
                 for num_args in range(start_index, end_index + 1):
+                    ret_suffix = ""
                     combinations = [
                         f"__ent{arg}: _TCCA[_T{arg}]"
                         for arg in range(num_args)
                     ]
+
+                    if num_args == end_index:
+                        ret_suffix = ", Unpack[TupleAny]"
+                        extra_args = (
+                            f", *entities: _ColumnsClauseArgument[Any]"
+                            f"{extra_args.replace(', *', '')}"
+                        )
+
                     buf.write(
                         textwrap.indent(
                             f"""
 @overload
 def {current_fnname}(
     {'self, ' if use_self else ''}{", ".join(combinations)},/{extra_args}
-) -> {return_type}[{', '.join(f'_T{i}' for i in range(num_args))}]:
+) -> {return_type}[{', '.join(f'_T{i}' for i in range(num_args))}{ret_suffix}]:
     ...
 
 """,  # noqa: E501