From 889164cb0c15f2ba9b57ad0154f0986bd22b9b87 Mon Sep 17 00:00:00 2001
From: Sebastian Rittau <srittau@rittau.biz>
Date: Tue, 8 Mar 2022 17:13:19 +0100
Subject: [PATCH] psycopg2: Annotate cursor execute* and dunder methods (#7460)

---
 stubs/psycopg2/psycopg2/_psycopg.pyi | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi
index 13aad3e7f..b4983f65d 100644
--- a/stubs/psycopg2/psycopg2/_psycopg.pyi
+++ b/stubs/psycopg2/psycopg2/_psycopg.pyi
@@ -1,8 +1,13 @@
-from typing import Any, Callable, TypeVar, overload
+from _typeshed import Self
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from types import TracebackType
+from typing import Any, TypeVar, overload
 
 import psycopg2
 import psycopg2.extensions
 
+_Vars = Sequence[Any] | Mapping[str, Any] | None
+
 BINARY: Any
 BINARYARRAY: Any
 BOOLEAN: Any
@@ -358,8 +363,8 @@ class cursor:
     def copy_expert(self, sql, file, size=...): ...
     def copy_from(self, file, table, sep=..., null=..., size=..., columns=...): ...
     def copy_to(self, file, table, sep=..., null=..., columns=...): ...
-    def execute(self, query, vars=...): ...
-    def executemany(self, query, vars_list): ...
+    def execute(self, query: str | bytes, vars: _Vars = ...) -> None: ...
+    def executemany(self, query: str | bytes, vars_list: Iterable[_Vars]) -> None: ...
     def fetchall(self) -> list[tuple[Any, ...]]: ...
     def fetchmany(self, size=...) -> list[tuple[Any, ...]]: ...
     def fetchone(self) -> tuple[Any, ...] | Any: ...
@@ -368,10 +373,12 @@ class cursor:
     def scroll(self, value, mode=...): ...
     def setinputsizes(self, sizes): ...
     def setoutputsize(self, size, column=...): ...
-    def __enter__(self): ...
-    def __exit__(self, type, value, traceback): ...
-    def __iter__(self): ...
-    def __next__(self): ...
+    def __enter__(self: Self) -> Self: ...
+    def __exit__(
+        self, type: type[BaseException] | None, value: BaseException | None, traceback: TracebackType | None
+    ) -> None: ...
+    def __iter__(self: Self) -> Self: ...
+    def __next__(self) -> tuple[Any, ...]: ...
 
 class lobject:
     closed: Any