improve rowmapping key type

CaselIT · CaselIT · commit b2ee1df06b13 · 2025-02-26T20:32:09.000+01:00
the accepted keys are also orm attributes, column elements, functions
etc, not only columns

Change-Id: I354de9b9668bc02b8b305a3c1f065744b28f8030
diff --git a/lib/sqlalchemy/engine/result.py b/lib/sqlalchemy/engine/result.py
@@ -51,11 +51,11 @@
 from ..util.typing import Unpack
 
 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
@@ -3442,7 +3442,7 @@ def _result_has_identity_key(self, result, adapter=None):
 
     def identity_key_from_row(
         self,
-        row: Optional[Union[Row[Unpack[TupleAny]], RowMapping]],
+        row: Union[Row[Unpack[TupleAny]], RowMapping],
         identity_token: Optional[Any] = None,
         adapter: Optional[ORMAdapter] = None,
     ) -> _IdentityKeyType[_O]:
@@ -3461,14 +3461,15 @@ def identity_key_from_row(
         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
@@ -8,6 +8,7 @@
 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
@@ -117,9 +118,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:
