From fc3293cf755dcfe08807d8e0832e0f6c18f916a3 Mon Sep 17 00:00:00 2001 From: Sebastian Rittau Date: Tue, 8 Mar 2022 15:41:56 +0100 Subject: [PATCH 1/3] psycopg2: Annotate cursor.execute(many) --- stubs/psycopg2/psycopg2/_psycopg.pyi | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi index 13aad3e7f8c0..20663ebc3f86 100644 --- a/stubs/psycopg2/psycopg2/_psycopg.pyi +++ b/stubs/psycopg2/psycopg2/_psycopg.pyi @@ -1,8 +1,11 @@ -from typing import Any, Callable, TypeVar, overload +from collections.abc import Callable, Iterable, Mapping, Sequence +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 +361,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: ... From a9c19ae29a89530369bdcef9937fad1d126ff812 Mon Sep 17 00:00:00 2001 From: Sebastian Rittau Date: Tue, 8 Mar 2022 16:10:11 +0100 Subject: [PATCH 2/3] Annotate cursor dunder methods --- stubs/psycopg2/psycopg2/_psycopg.pyi | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi index 20663ebc3f86..123545868d5c 100644 --- a/stubs/psycopg2/psycopg2/_psycopg.pyi +++ b/stubs/psycopg2/psycopg2/_psycopg.pyi @@ -1,4 +1,6 @@ +from _typeshed import Self from collections.abc import Callable, Iterable, Mapping, Sequence +from types import TracebackType from typing import Any, TypeVar, overload import psycopg2 @@ -371,10 +373,10 @@ 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: ... + def __iter__(self: Self) -> Self: ... + def __next__(self) -> tuple[Any, ...]: ... class lobject: closed: Any From 1e1324c7654c795ab8c3c3ca4ce1c838f2bae8fc Mon Sep 17 00:00:00 2001 From: Sebastian Rittau Date: Tue, 8 Mar 2022 16:24:12 +0100 Subject: [PATCH 3/3] Add missing None --- stubs/psycopg2/psycopg2/_psycopg.pyi | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi index 123545868d5c..b4983f65d2a0 100644 --- a/stubs/psycopg2/psycopg2/_psycopg.pyi +++ b/stubs/psycopg2/psycopg2/_psycopg.pyi @@ -374,7 +374,9 @@ class cursor: def setinputsizes(self, sizes): ... def setoutputsize(self, size, column=...): ... def __enter__(self: Self) -> Self: ... - def __exit__(self, type: type[BaseException] | None, value: BaseException | None, traceback: TracebackType) -> None: ... + def __exit__( + self, type: type[BaseException] | None, value: BaseException | None, traceback: TracebackType | None + ) -> None: ... def __iter__(self: Self) -> Self: ... def __next__(self) -> tuple[Any, ...]: ...