psycopg2: Add missing modules, add annotations (#10630)

stubs/psycopg2/@tests/stubtest_allowlist.txt

@@ -1,5 +1,3 @@
-psycopg2.connection
-psycopg2.cursor
-psycopg2.pool.AbstractConnectionPool.closeall
-psycopg2.pool.AbstractConnectionPool.getconn
-psycopg2.pool.AbstractConnectionPool.putconn
+psycopg2.pool.AbstractConnectionPool.(closeall|getconn|putconn)
+psycopg2.(extras|_range).RangeAdapter.name
+psycopg2.(_psycopg|extensions).connection.async  # async is a reserved keyword in Python 3.7

stubs/psycopg2/@tests/test_cases/check_connect.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing_extensions import assert_type
+
+import psycopg2
+from psycopg2.extensions import connection, cursor
+
+
+class MyCursor(cursor):
+    pass
+
+
+class MyConnection(connection):
+    pass
+
+
+def custom_connection(dsn: str) -> MyConnection:
+    return MyConnection(dsn, async_=0)
+
+
+# -> psycopg2.extensions.connection
+assert_type(psycopg2.connect(), connection)
+assert_type(psycopg2.connect("test-conn"), connection)
+assert_type(psycopg2.connect(None), connection)
+assert_type(psycopg2.connect("test-conn", connection_factory=None), connection)
+
+assert_type(psycopg2.connect(cursor_factory=MyCursor), connection)
+assert_type(psycopg2.connect("test-conn", cursor_factory=MyCursor), connection)
+assert_type(psycopg2.connect(None, cursor_factory=MyCursor), connection)
+assert_type(psycopg2.connect("test-conn", connection_factory=None, cursor_factory=MyCursor), connection)
+
+# -> custom_connection
+assert_type(psycopg2.connect(connection_factory=MyConnection), MyConnection)
+assert_type(psycopg2.connect("test-conn", connection_factory=MyConnection), MyConnection)
+assert_type(psycopg2.connect("test-conn", MyConnection), MyConnection)
+assert_type(psycopg2.connect(connection_factory=custom_connection), MyConnection)
+
+assert_type(psycopg2.connect(connection_factory=MyConnection, cursor_factory=MyCursor), MyConnection)
+assert_type(psycopg2.connect("test-conn", connection_factory=MyConnection, cursor_factory=MyCursor), MyConnection)
+assert_type(psycopg2.connect(connection_factory=custom_connection, cursor_factory=MyCursor), MyConnection)

stubs/psycopg2/@tests/test_cases/check_extensions.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing_extensions import assert_type
+
+import psycopg2.extensions
+import psycopg2.extras
+from psycopg2.extensions import make_dsn
+
+# make_dsn
+# --------
+
+# (None) -> str
+assert_type(make_dsn(), str)
+assert_type(make_dsn(None), str)
+assert_type(make_dsn(dsn=None), str)
+
+# (bytes) -> bytes
+assert_type(make_dsn(b""), bytes)
+assert_type(make_dsn(dsn=b""), bytes)
+
+# (bytes, **Kwargs) -> str
+assert_type(make_dsn(b"", database=""), str)
+assert_type(make_dsn(dsn=b"", database=""), str)
+
+# (str, **OptionalKwargs) -> str
+assert_type(make_dsn(""), str)
+assert_type(make_dsn(dsn=""), str)
+assert_type(make_dsn("", database=None), str)
+assert_type(make_dsn(dsn="", database=None), str)
+
+
+# connection.cursor
+# -----------------
+
+# (name?, None?, ...) -> psycopg2.extensions.cursor
+conn = psycopg2.connect("test-conn")
+assert_type(conn.cursor(), psycopg2.extensions.cursor)
+assert_type(conn.cursor("test-cur"), psycopg2.extensions.cursor)
+assert_type(conn.cursor("test-cur", None), psycopg2.extensions.cursor)
+assert_type(conn.cursor("test-cur", cursor_factory=None), psycopg2.extensions.cursor)
+
+
+# (name?, cursor_factory(), ...) -> custom_cursor
+class MyCursor(psycopg2.extensions.cursor):
+    pass
+
+
+assert_type(conn.cursor("test-cur", cursor_factory=MyCursor), MyCursor)
+assert_type(conn.cursor("test-cur", cursor_factory=lambda c, n: MyCursor(c, n)), MyCursor)
+
+dconn = psycopg2.extras.DictConnection("test-dconn")
+assert_type(dconn.cursor(), psycopg2.extras.DictCursor)
+assert_type(dconn.cursor("test-dcur"), psycopg2.extras.DictCursor)
+assert_type(dconn.cursor("test-dcur", None), psycopg2.extras.DictCursor)
+assert_type(dconn.cursor("test-dcur", cursor_factory=None), psycopg2.extras.DictCursor)
+assert_type(dconn.cursor("test-dcur", cursor_factory=MyCursor), MyCursor)

stubs/psycopg2/METADATA.toml

@@ -3,4 +3,4 @@ upstream_repository = "https://github.com/psycopg/psycopg2"
 partial_stub = true
 
 [tool.stubtest]
-ignore_missing_stub = true
+ignore_missing_stub = false

stubs/psycopg2/psycopg2/__init__.pyi

@@ -27,19 +27,28 @@ from psycopg2._psycopg import (
     Warning as Warning,
     __libpq_version__ as __libpq_version__,
     apilevel as apilevel,
-    connection as connection,
-    cursor as cursor,
+    connection,
+    cursor,
     paramstyle as paramstyle,
     threadsafety as threadsafety,
 )
 
 _T_conn = TypeVar("_T_conn", bound=connection)
 
 @overload
-def connect(dsn: str, connection_factory: Callable[..., _T_conn], cursor_factory: None = None, **kwargs: Any) -> _T_conn: ...
+def connect(
+    dsn: str | None,
+    connection_factory: Callable[..., _T_conn],
+    cursor_factory: Callable[..., cursor] | None = None,
+    **kwargs: Any,
+) -> _T_conn: ...
 @overload
 def connect(
-    dsn: str | None = None, *, connection_factory: Callable[..., _T_conn], cursor_factory: None = None, **kwargs: Any
+    dsn: str | None = None,
+    *,
+    connection_factory: Callable[..., _T_conn],
+    cursor_factory: Callable[..., cursor] | None = None,
+    **kwargs: Any,
 ) -> _T_conn: ...
 @overload
 def connect(

stubs/psycopg2/psycopg2/_ipaddress.pyi

@@ -1,9 +1,9 @@
-from _typeshed import Incomplete
-from typing import Any
+import ipaddress as ipaddress
+from _typeshed import Unused
 
-ipaddress: Any
+from psycopg2._psycopg import QuotedString, connection, cursor
 
-def register_ipaddress(conn_or_curs: Incomplete | None = None) -> None: ...
-def cast_interface(s, cur: Incomplete | None = None): ...
-def cast_network(s, cur: Incomplete | None = None): ...
-def adapt_ipaddress(obj): ...
+def register_ipaddress(conn_or_curs: connection | cursor | None = None) -> None: ...
+def cast_interface(s: str, cur: Unused = None) -> ipaddress.IPv4Interface | ipaddress.IPv6Interface | None: ...
+def cast_network(s: str, cur: Unused = None) -> ipaddress.IPv4Network | ipaddress.IPv6Network | None: ...
+def adapt_ipaddress(obj: object) -> QuotedString: ...

stubs/psycopg2/psycopg2/_json.pyi

@@ -1,5 +1,8 @@
-from _typeshed import Incomplete
+from collections.abc import Callable
 from typing import Any
+from typing_extensions import Self
+
+from psycopg2._psycopg import _type, connection, cursor
 
 JSON_OID: int
 JSONARRAY_OID: int
@@ -8,19 +11,23 @@ JSONBARRAY_OID: int
 
 class Json:
     adapted: Any
-    def __init__(self, adapted, dumps: Incomplete | None = None) -> None: ...
-    def __conform__(self, proto): ...
-    def dumps(self, obj): ...
-    def prepare(self, conn) -> None: ...
-    def getquoted(self): ...
+    def __init__(self, adapted: Any, dumps: Callable[..., str] | None = None) -> None: ...
+    def __conform__(self, proto) -> Self | None: ...
+    def dumps(self, obj: Any) -> str: ...
+    def prepare(self, conn: connection | None) -> None: ...
+    def getquoted(self) -> bytes: ...
 
 def register_json(
-    conn_or_curs: Incomplete | None = None,
+    conn_or_curs: connection | cursor | None = None,
     globally: bool = False,
-    loads: Incomplete | None = None,
-    oid: Incomplete | None = None,
-    array_oid: Incomplete | None = None,
+    loads: Callable[..., Any] | None = None,
+    oid: int | None = None,
+    array_oid: int | None = None,
     name: str = "json",
-): ...
-def register_default_json(conn_or_curs: Incomplete | None = None, globally: bool = False, loads: Incomplete | None = None): ...
-def register_default_jsonb(conn_or_curs: Incomplete | None = None, globally: bool = False, loads: Incomplete | None = None): ...
+) -> tuple[_type, _type | None]: ...
+def register_default_json(
+    conn_or_curs: connection | cursor | None = None, globally: bool = False, loads: Callable[..., Any] | None = None
+) -> tuple[_type, _type | None]: ...
+def register_default_jsonb(
+    conn_or_curs: connection | cursor | None = None, globally: bool = False, loads: Callable[..., Any] | None = None
+) -> tuple[_type, _type | None]: ...