Improve from_select typing

Properly type :paramref:`_dml.Insert.from_select.names` to accept
a list of string or columns or mapped attributes.
Fixes: #9514
Change-Id: I2a3396454e79bd2113fc3d57d1252a2e6d141310

diff --git a/doc/build/changelog/unreleased_20/9514.rst b/doc/build/changelog/unreleased_20/9514.rst
new file mode 100644 (file)
index 0000000..cb2d776
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/9514.rst
@@ -0,0 +1,6 @@
+.. change::
+    :tags: bug, orm, typing
+    :tickets: 9514
+
+    Properly type :paramref:`_dml.Insert.from_select.names` to accept
+    a list of string or columns or mapped attributes.

diff --git a/lib/sqlalchemy/sql/dml.py b/lib/sqlalchemy/sql/dml.py
index dbbf09f1b930caef5eb7da878c4fc28ae20717ec..316392c0cd2f529e57406e253ca68ed1ceaa0435 100644 (file)
--- a/lib/sqlalchemy/sql/dml.py
+++ b/lib/sqlalchemy/sql/dml.py
@@ -228,8 +228,7 @@ class DMLState(CompileState):
     def _process_select_values(self, statement: ValuesBase) -> None:
         assert statement._select_names is not None
         parameters: MutableMapping[_DMLColumnElement, Any] = {
-            coercions.expect(roles.DMLColumnRole, name, as_key=True): Null()
-            for name in statement._select_names
+            name: Null() for name in statement._select_names
         }
 
         if self._no_parameters:
@@ -1195,7 +1194,7 @@ class Insert(ValuesBase):
     @_generative
     def from_select(
         self,
-        names: List[str],
+        names: Sequence[_DMLColumnArgument],
         select: Selectable,
         include_defaults: bool = True,
     ) -> Self:
@@ -1253,7 +1252,10 @@ class Insert(ValuesBase):
                 "This construct already inserts value expressions"
             )
 
-        self._select_names = names
+        self._select_names = [
+            coercions.expect(roles.DMLColumnRole, name, as_key=True)
+            for name in names
+        ]
         self._inline = True
         self.include_insert_from_select_defaults = include_defaults
         self.select = coercions.expect(roles.DMLSelectRole, select)

diff --git a/test/ext/mypy/plain_files/dml.py b/test/ext/mypy/plain_files/dml.py
index d2ffbf1e1651fc304b36b135f9197bde1e2709fc..596f47795c41b7330a09c3f2819cb2e776990f74 100644 (file)
--- a/test/ext/mypy/plain_files/dml.py
+++ b/test/ext/mypy/plain_files/dml.py
@@ -5,6 +5,7 @@ from typing import Dict
 
 from sqlalchemy import Column
 from sqlalchemy import insert
+from sqlalchemy import select
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import Mapped
 from sqlalchemy.orm import mapped_column
@@ -33,3 +34,8 @@ stmt2 = insert(User).values(d2)
 
 d3: Dict[Column[str], Any] = {}
 stmt3 = insert(User).values(d3)
+
+stmt4 = insert(User).from_select(
+    [User.id, "name", User.__table__.c.data],
+    select(User.id, User.name, User.data),
+)