From 2dc24ced0e94c50623aab41c75e42883d68a846a Mon Sep 17 00:00:00 2001
From: Federico Caselli <cfederico87@gmail.com>
Date: Mon, 27 Mar 2023 22:03:44 +0200
Subject: [PATCH] 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
---
 doc/build/changelog/unreleased_20/9514.rst |  6 ++++++
 lib/sqlalchemy/sql/dml.py                  | 10 ++++++----
 test/ext/mypy/plain_files/dml.py           |  6 ++++++
 3 files changed, 18 insertions(+), 4 deletions(-)
 create mode 100644 doc/build/changelog/unreleased_20/9514.rst

diff --git a/doc/build/changelog/unreleased_20/9514.rst b/doc/build/changelog/unreleased_20/9514.rst
new file mode 100644
index 0000000000..cb2d7761ee
--- /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 dbbf09f1b9..316392c0cd 100644
--- 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 d2ffbf1e16..596f47795c 100644
--- 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),
+)
-- 
2.47.2