From: Federico Caselli <cfederico87@gmail.com>
Date: Tue, 25 Feb 2025 22:06:55 +0000 (+0100)
Subject: improve rowmapping key type
X-Git-Tag: rel_2_0_39~10^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=b6251cb996c59723a2580d2672d9cd4070048c69;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

improve rowmapping key type

the accepted keys are also orm attributes, column elements, functions
etc, not only columns

Change-Id: I354de9b9668bc02b8b305a3c1f065744b28f8030
(cherry picked from commit b2ee1df06b138fc9588ea312d4a477699ec9b5d0)
---

diff --git a/lib/sqlalchemy/engine/result.py b/lib/sqlalchemy/engine/result.py
index 7411fd74f6..3c81fc6052 100644
--- a/lib/sqlalchemy/engine/result.py
+++ b/lib/sqlalchemy/engine/result.py
@@ -52,11 +52,11 @@ else:
     from sqlalchemy.cyextension.resultproxy import tuplegetter as tuplegetter
 
 if typing.TYPE_CHECKING:
-    from ..sql.schema import Column
+    from ..sql.elements import SQLCoreOperations
     from ..sql.type_api import _ResultProcessorType
 
-_KeyType = Union[str, "Column[Any]"]
-_KeyIndexType = Union[str, "Column[Any]", int]
+_KeyType = Union[str, "SQLCoreOperations[Any]"]
+_KeyIndexType = Union[_KeyType, int]
 
 # is overridden in cursor using _CursorKeyMapRecType
 _KeyMapRecType = Any
diff --git a/lib/sqlalchemy/orm/mapper.py b/lib/sqlalchemy/orm/mapper.py
index 11010efbcd..eab2be558f 100644
--- a/lib/sqlalchemy/orm/mapper.py
+++ b/lib/sqlalchemy/orm/mapper.py
@@ -3440,7 +3440,7 @@ class Mapper(
 
     def identity_key_from_row(
         self,
-        row: Optional[Union[Row[Any], RowMapping]],
+        row: Union[Row[Any], RowMapping],
         identity_token: Optional[Any] = None,
         adapter: Optional[ORMAdapter] = None,
     ) -> _IdentityKeyType[_O]:
@@ -3459,14 +3459,15 @@ class Mapper(
         if adapter:
             pk_cols = [adapter.columns[c] for c in pk_cols]
 
+        mapping: RowMapping
         if hasattr(row, "_mapping"):
-            mapping = row._mapping  # type: ignore
+            mapping = row._mapping
         else:
-            mapping = cast("Mapping[Any, Any]", row)
+            mapping = row  # type: ignore[assignment]
 
         return (
             self._identity_class,
-            tuple(mapping[column] for column in pk_cols),  # type: ignore
+            tuple(mapping[column] for column in pk_cols),
             identity_token,
         )
 
diff --git a/test/typing/plain_files/sql/typed_results.py b/test/typing/plain_files/sql/typed_results.py
index c7842a7e79..9ed591815a 100644
--- a/test/typing/plain_files/sql/typed_results.py
+++ b/test/typing/plain_files/sql/typed_results.py
@@ -9,6 +9,7 @@ from typing import Type
 from sqlalchemy import Column
 from sqlalchemy import column
 from sqlalchemy import create_engine
+from sqlalchemy import func
 from sqlalchemy import insert
 from sqlalchemy import Integer
 from sqlalchemy import MetaData
@@ -118,9 +119,22 @@ def t_result_ctxmanager() -> None:
         reveal_type(r4)
 
 
-def t_core_mappings() -> None:
+def t_mappings() -> None:
     r = connection.execute(select(t_user)).mappings().one()
-    r.get(t_user.c.id)
+    r["name"]  # string
+    r.get(t_user.c.id)  # column
+
+    r2 = connection.execute(select(User)).mappings().one()
+    r2[User.id]  # orm attribute
+    r2[User.__table__.c.id]  # form clause column
+
+    m2 = User.id * 2
+    s2 = User.__table__.c.id + 2
+    fn = func.abs(User.id)
+    r3 = connection.execute(select(m2, s2, fn)).mappings().one()
+    r3[m2]  # col element
+    r3[s2]  # also col element
+    r3[fn]  # function
 
 
 def t_entity_varieties() -> None: