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)