From ae39348d7a7397542ec051601a576a09f0345069 Mon Sep 17 00:00:00 2001
From: Sebastian Rittau <srittau@rittau.biz>
Date: Thu, 17 Feb 2022 18:27:55 +0100
Subject: [PATCH] Annotate Session.execute() and Result (#7252)

* Annotate return type of Session.execute()

* Annotate Result return types

* Remove sub-class overrides with identical signatures
---
 stubs/SQLAlchemy/sqlalchemy/engine/cursor.pyi |  1 -
 stubs/SQLAlchemy/sqlalchemy/engine/result.pyi | 29 ++++++++++---------
 stubs/SQLAlchemy/sqlalchemy/orm/session.pyi   |  3 +-
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/stubs/SQLAlchemy/sqlalchemy/engine/cursor.pyi b/stubs/SQLAlchemy/sqlalchemy/engine/cursor.pyi
index 557743fb0..980290818 100644
--- a/stubs/SQLAlchemy/sqlalchemy/engine/cursor.pyi
+++ b/stubs/SQLAlchemy/sqlalchemy/engine/cursor.pyi
@@ -115,7 +115,6 @@ class BaseCursorResult:
 class CursorResult(BaseCursorResult, Result):
     def merge(self, *others): ...
     def close(self) -> None: ...
-    def yield_per(self, num) -> None: ...
 
 class LegacyCursorResult(CursorResult):
     def close(self) -> None: ...
diff --git a/stubs/SQLAlchemy/sqlalchemy/engine/result.pyi b/stubs/SQLAlchemy/sqlalchemy/engine/result.pyi
index 8aab5bc9d..e6bb880db 100644
--- a/stubs/SQLAlchemy/sqlalchemy/engine/result.pyi
+++ b/stubs/SQLAlchemy/sqlalchemy/engine/result.pyi
@@ -1,7 +1,9 @@
+from _typeshed import Self
 from collections.abc import Generator, KeysView
 from typing import Any
 
 from ..sql.base import InPlaceGenerative
+from .row import Row
 
 class ResultMetaData:
     @property
@@ -36,24 +38,24 @@ class _WithKeys:
 class Result(_WithKeys, ResultInternal):
     def __init__(self, cursor_metadata) -> None: ...
     def close(self) -> None: ...
-    def yield_per(self, num: int) -> None: ...
-    def unique(self, strategy: Any | None = ...) -> None: ...
+    def yield_per(self: Self, num: int) -> Self: ...
+    def unique(self: Self, strategy: Any | None = ...) -> Self: ...
     def columns(self, *col_expressions): ...
     def scalars(self, index: int = ...) -> ScalarResult: ...
     def mappings(self) -> MappingResult: ...
     def __iter__(self): ...
     def __next__(self): ...
-    def partitions(self, size: int | None = ...) -> Generator[Any, None, None]: ...
-    def fetchall(self) -> list[Any]: ...
-    def fetchone(self) -> Any | None: ...
-    def fetchmany(self, size: int | None = ...) -> list[Any]: ...
-    def all(self) -> list[Any]: ...
-    def first(self): ...
-    def one_or_none(self): ...
-    def scalar_one(self): ...
-    def scalar_one_or_none(self): ...
-    def one(self): ...
-    def scalar(self): ...
+    def partitions(self, size: int | None = ...) -> Generator[list[Row], None, None]: ...
+    def fetchall(self) -> list[Row]: ...
+    def fetchone(self) -> Row | None: ...
+    def fetchmany(self, size: int | None = ...) -> list[Row]: ...
+    def all(self) -> list[Row]: ...
+    def first(self) -> Row | None: ...
+    def one_or_none(self) -> Row | None: ...
+    def scalar_one(self) -> Any: ...
+    def scalar_one_or_none(self) -> Any | None: ...
+    def one(self) -> Row: ...
+    def scalar(self) -> Any | None: ...
     def freeze(self) -> FrozenResult: ...
     def merge(self, *others) -> MergedResult: ...
 
@@ -110,7 +112,6 @@ class ChunkedIteratorResult(IteratorResult):
     def __init__(
         self, cursor_metadata, chunks, source_supports_scalars: bool = ..., raw: Any | None = ..., dynamic_yield_per: bool = ...
     ) -> None: ...
-    def yield_per(self, num) -> None: ...
 
 class MergedResult(IteratorResult):
     closed: bool
diff --git a/stubs/SQLAlchemy/sqlalchemy/orm/session.pyi b/stubs/SQLAlchemy/sqlalchemy/orm/session.pyi
index 18aa5a589..68586fdf9 100644
--- a/stubs/SQLAlchemy/sqlalchemy/orm/session.pyi
+++ b/stubs/SQLAlchemy/sqlalchemy/orm/session.pyi
@@ -2,6 +2,7 @@ from collections.abc import Mapping
 from typing import Any, TypeVar, overload
 
 from ..engine.base import Connection
+from ..engine.result import Result
 from ..engine.util import TransactionalContext
 from ..sql.elements import ColumnElement
 from ..sql.schema import Table
@@ -135,7 +136,7 @@ class Session(_SessionClassMethods):
         _parent_execute_state: Any | None = ...,
         _add_event: Any | None = ...,
         **kw,
-    ): ...
+    ) -> Result: ...
     def scalar(self, statement, params: Any | None = ..., execution_options=..., bind_arguments: Any | None = ..., **kw): ...
     def scalars(self, statement, params: Any | None = ..., execution_options=..., bind_arguments: Any | None = ..., **kw): ...
     def close(self) -> None: ...