Add typing overloads to Query.__getitem__ and AppenderQuery.__getitem__

Fixes #13128

### Description
`Query.__getitem__` and `AppenderQuery.__getitem__` previously returned Union[_T, List[_T]] for all inputs, making the return type inaccurate.
Added `@overload` signatures so that integer index returns _T and slice returns List[_T].

This pull request is:

- [x] A documentation / typographical / small typing error fix
	- Good to go, no issue or tests are needed

Closes: #13142
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/13142
Pull-request-sha: 9ba1f0145d90b18c997137aeba7fde72dac23a7c

Change-Id: Ib37ab63d3d844491c34cc5ccfc4efc1591a1878c

diff --git a/lib/sqlalchemy/orm/dynamic.py b/lib/sqlalchemy/orm/dynamic.py
index dfa22dd5286605ce225f662b88dd42fa1d76eace..b7ed1c9399d88c72e17290dfc2ac483386aaa650 100644 (file)
--- a/lib/sqlalchemy/orm/dynamic.py
+++ b/lib/sqlalchemy/orm/dynamic.py
@@ -24,6 +24,7 @@ from typing import Iterable
 from typing import Iterator
 from typing import List
 from typing import Optional
+from typing import overload
 from typing import Tuple
 from typing import Type
 from typing import TYPE_CHECKING
@@ -176,7 +177,13 @@ class _AppenderMixin(_AbstractCollectionWriter[_T]):
 
         def __iter__(self) -> Iterator[_T]: ...
 
-    def __getitem__(self, index: Any) -> Union[_T, List[_T]]:
+    @overload
+    def __getitem__(self, index: int) -> _T: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> List[_T]: ...
+
+    def __getitem__(self, index: Union[int, slice]) -> Union[_T, List[_T]]:
         sess = self.session
         if sess is None:
             return self.attr._get_collection_history(
@@ -184,7 +191,7 @@ class _AppenderMixin(_AbstractCollectionWriter[_T]):
                 PassiveFlag.PASSIVE_NO_INITIALIZE,
             ).indexed(index)
         else:
-            return self._generate(sess).__getitem__(index)  # type: ignore[no-any-return] # noqa: E501
+            return self._generate(sess).__getitem__(index)
 
     def count(self) -> int:
         sess = self.session

diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index a4f94785140bcb07d13c91a612d98c7236070498..38ee0297c3755953a54eb2fc83147f01ab71ecdc 100644 (file)
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -35,6 +35,7 @@ from typing import Mapping
 from typing import Optional
 from typing import overload
 from typing import Sequence
+from typing import SupportsIndex
 from typing import Tuple
 from typing import Type
 from typing import TYPE_CHECKING
@@ -2588,6 +2589,12 @@ class Query(
         self._set_select_from(from_obj, False)
         return self
 
+    @overload
+    def __getitem__(self, item: slice) -> List[_T]: ...
+
+    @overload
+    def __getitem__(self, item: SupportsIndex) -> _T: ...
+
     def __getitem__(self, item: Any) -> Any:
         return orm_util._getitem(
             self,

diff --git a/test/typing/plain_files/orm/dynamic_rel.py b/test/typing/plain_files/orm/dynamic_rel.py
index c9ebdbc1e441038ea4531544c755718c8749fa44..4bcc6addbcf48eba8b65e6f7006b6bb7060c6bcf 100644 (file)
--- a/test/typing/plain_files/orm/dynamic_rel.py
+++ b/test/typing/plain_files/orm/dynamic_rel.py
@@ -80,3 +80,8 @@ with Session() as session:
 
     # test #9985
     stmt = select(User).join(User.addresses)
+
+    # test #13128
+    if typing.TYPE_CHECKING:
+        assert_type(u.addresses[0], Address)
+        assert_type(u.addresses[1:3], list[Address])

diff --git a/test/typing/plain_files/orm/typed_queries.py b/test/typing/plain_files/orm/typed_queries.py
index da901ab752cdb9fe736ee941dce7aafc62410f66..9b15d29e8e40a6ece47cb399845760124f764f61 100644 (file)
--- a/test/typing/plain_files/orm/typed_queries.py
+++ b/test/typing/plain_files/orm/typed_queries.py
@@ -512,3 +512,9 @@ def t_aliased_fromclause() -> None:
 def test_select_from() -> None:
     select(1).select_from(User).exists()
     exists(1).select_from(User).select()
+
+
+def t_legacy_query_getitem() -> None:
+    q1 = session.query(User).filter(User.id == 5)
+    assert_type(q1[0], User)
+    assert_type(q1[1:3], list[User])