Additional return types for psycopg2 connections (#8528)

* Return types for psycopg2 dict connections

DB connection and cursor methods that deal with `dict`-style results now indicate the expected return types as implemented.

* Return types for psycopg2 namedtuple connections

DB connection and cursor methods that deal with `namedtuple` results now indicate the expected return types as implemented.

* Use ABC iterator

As required by flake8, the `Iterator` referenced in psycopg2's `extras` module has been switched to the `collections.abc.Iterator[T]` variant.

* Fix base psycopg2 cursor iter return type

The previous return type of `Self` is wrong; `cursor`'s `__iter__` method returns an `Iterator`.

* Correct return type for cursor iter and next methods

The previous attempt to fix the return type of `cursor`'s (and subclasses) `__iter__` method was misguided; they should indeed return `Self`. It's the `__next__` methods that had to be updated to return the correct record/row instance type.

* Comprehensive overloads for psycopg2 extra connections

Provides full method signatures for the `cursor` methods of the `DictConnection`, `RealDictConnection` and `NamedTupleConnection` types. Notably this includes a default value for `cursor_factory` in each case, while preserving the option to override the parameter manually.

* Have mypy ignore incompatible psycopg2 return types

The return types of the `fetch*` and `__next__` methods of `DictCursor`, `RealDictCursor` and `NamedTupleCursor` are incompatible with the base `cursor` class's corresponding methods' return types. However, this does accurately reflect reality, so ignore the mypy errors in those cases.

* Use ABC callable for psycopg2 extras module

As required by flake8, the `Callable` referenced in psycopg2's `extras` module has been switched to the `Callable` variant.

* Remove superfluous psycopg2 member overrides

Several members in the `DictCursorBase` type definition were entirely unnecessary, so they have been removed. In addition, adds a type to the `size` param of the `cursor`'s `fetchmany` method.

stubs/psycopg2/psycopg2/_psycopg.pyi

@@ -98,7 +98,7 @@ class cursor:
     def execute(self, query: str | bytes | Composable, vars: _Vars = ...) -> None: ...
     def executemany(self, query: str | bytes | Composable, vars_list: Iterable[_Vars]) -> None: ...
     def fetchall(self) -> list[tuple[Any, ...]]: ...
-    def fetchmany(self, size=...) -> list[tuple[Any, ...]]: ...
+    def fetchmany(self, size: int | None = ...) -> list[tuple[Any, ...]]: ...
     def fetchone(self) -> tuple[Any, ...] | None: ...
     def mogrify(self, *args, **kwargs): ...
     def nextset(self): ...

stubs/psycopg2/psycopg2/extras.pyi

@@ -1,5 +1,6 @@
 from collections import OrderedDict
-from typing import Any
+from collections.abc import Callable
+from typing import Any, NamedTuple, TypeVar, overload
 
 from psycopg2._ipaddress import register_ipaddress as register_ipaddress
 from psycopg2._json import (
@@ -28,22 +29,37 @@ from psycopg2._range import (
 
 from .extensions import connection as _connection, cursor as _cursor, quote_ident as quote_ident
 
+_T_cur = TypeVar("_T_cur", bound=_cursor)
+
 class DictCursorBase(_cursor):
-    row_factory: Any
     def __init__(self, *args, **kwargs) -> None: ...
-    def fetchone(self) -> tuple[Any, ...] | None: ...
-    def fetchmany(self, size: int | None = ...) -> list[tuple[Any, ...]]: ...
-    def fetchall(self) -> list[tuple[Any, ...]]: ...
-    def __iter__(self): ...
 
 class DictConnection(_connection):
-    def cursor(self, *args, **kwargs): ...
+    @overload
+    def cursor(self, name: str | bytes | None = ..., *, withhold: bool = ..., scrollable: bool | None = ...) -> DictCursor: ...
+    @overload
+    def cursor(
+        self,
+        name: str | bytes | None = ...,
+        *,
+        cursor_factory: Callable[..., _T_cur],
+        withhold: bool = ...,
+        scrollable: bool | None = ...,
+    ) -> _T_cur: ...
+    @overload
+    def cursor(
+        self, name: str | bytes | None, cursor_factory: Callable[..., _T_cur], withhold: bool = ..., scrollable: bool | None = ...
+    ) -> _T_cur: ...
 
 class DictCursor(DictCursorBase):
     def __init__(self, *args, **kwargs) -> None: ...
     index: Any
     def execute(self, query, vars: Any | None = ...): ...
     def callproc(self, procname, vars: Any | None = ...): ...
+    def fetchone(self) -> DictRow | None: ...  # type: ignore[override]
+    def fetchmany(self, size: int | None = ...) -> list[DictRow]: ...  # type: ignore[override]
+    def fetchall(self) -> list[DictRow]: ...  # type: ignore[override]
+    def __next__(self) -> DictRow: ...  # type: ignore[override]
 
 class DictRow(list[Any]):
     def __init__(self, cursor) -> None: ...
@@ -58,20 +74,56 @@ class DictRow(list[Any]):
     def __reduce__(self): ...
 
 class RealDictConnection(_connection):
-    def cursor(self, *args, **kwargs): ...
+    @overload
+    def cursor(
+        self, name: str | bytes | None = ..., *, withhold: bool = ..., scrollable: bool | None = ...
+    ) -> RealDictCursor: ...
+    @overload
+    def cursor(
+        self,
+        name: str | bytes | None = ...,
+        *,
+        cursor_factory: Callable[..., _T_cur],
+        withhold: bool = ...,
+        scrollable: bool | None = ...,
+    ) -> _T_cur: ...
+    @overload
+    def cursor(
+        self, name: str | bytes | None, cursor_factory: Callable[..., _T_cur], withhold: bool = ..., scrollable: bool | None = ...
+    ) -> _T_cur: ...
 
 class RealDictCursor(DictCursorBase):
     def __init__(self, *args, **kwargs) -> None: ...
     column_mapping: Any
     def execute(self, query, vars: Any | None = ...): ...
     def callproc(self, procname, vars: Any | None = ...): ...
+    def fetchone(self) -> RealDictRow | None: ...  # type: ignore[override]
+    def fetchmany(self, size: int | None = ...) -> list[RealDictRow]: ...  # type: ignore[override]
+    def fetchall(self) -> list[RealDictRow]: ...  # type: ignore[override]
+    def __next__(self) -> RealDictRow: ...  # type: ignore[override]
 
 class RealDictRow(OrderedDict[Any, Any]):
     def __init__(self, *args, **kwargs) -> None: ...
     def __setitem__(self, key, value) -> None: ...
 
 class NamedTupleConnection(_connection):
-    def cursor(self, *args, **kwargs): ...
+    @overload
+    def cursor(
+        self, name: str | bytes | None = ..., *, withhold: bool = ..., scrollable: bool | None = ...
+    ) -> NamedTupleCursor: ...
+    @overload
+    def cursor(
+        self,
+        name: str | bytes | None = ...,
+        *,
+        cursor_factory: Callable[..., _T_cur],
+        withhold: bool = ...,
+        scrollable: bool | None = ...,
+    ) -> _T_cur: ...
+    @overload
+    def cursor(
+        self, name: str | bytes | None, cursor_factory: Callable[..., _T_cur], withhold: bool = ..., scrollable: bool | None = ...
+    ) -> _T_cur: ...
 
 class NamedTupleCursor(_cursor):
     Record: Any
@@ -79,10 +131,10 @@ class NamedTupleCursor(_cursor):
     def execute(self, query, vars: Any | None = ...): ...
     def executemany(self, query, vars): ...
     def callproc(self, procname, vars: Any | None = ...): ...
-    def fetchone(self) -> tuple[Any, ...] | None: ...
-    def fetchmany(self, size: int | None = ...) -> list[tuple[Any, ...]]: ...
-    def fetchall(self) -> list[tuple[Any, ...]]: ...
-    def __iter__(self): ...
+    def fetchone(self) -> NamedTuple | None: ...  # type: ignore[override]
+    def fetchmany(self, size: int | None = ...) -> list[NamedTuple]: ...  # type: ignore[override]
+    def fetchall(self) -> list[NamedTuple]: ...  # type: ignore[override]
+    def __next__(self) -> NamedTuple: ...  # type: ignore[override]
 
 class LoggingConnection(_connection):
     log: Any