From 8715951fca094fa801c46f860e9aba1b5018874a Mon Sep 17 00:00:00 2001
From: Joel <oldsneerjaw@gmail.com>
Date: Sat, 20 Aug 2022 16:34:13 -0700
Subject: [PATCH] Improve return type for psycopg2 connect function (#8567)

When a `connection_factory` argument is provided to psycopg2's `connect` function, the function's return type now matches that of the factory output class. However, if `cursor_factory` is set and has a non-`None` value and/or `connection_factory` is not set or is `None`, the return type is simply `connection`, as before.
---
 stubs/psycopg2/psycopg2/__init__.pyi | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/stubs/psycopg2/psycopg2/__init__.pyi b/stubs/psycopg2/psycopg2/__init__.pyi
index d1165cd2d..5c0277751 100644
--- a/stubs/psycopg2/psycopg2/__init__.pyi
+++ b/stubs/psycopg2/psycopg2/__init__.pyi
@@ -1,4 +1,5 @@
-from typing import Any
+from collections.abc import Callable
+from typing import Any, TypeVar, overload
 
 # connection and cursor not available at runtime
 from psycopg2._psycopg import (
@@ -32,6 +33,18 @@ from psycopg2._psycopg import (
     threadsafety as threadsafety,
 )
 
+_T_conn = TypeVar("_T_conn", bound=connection)
+
+@overload
+def connect(dsn: str, connection_factory: Callable[..., _T_conn], cursor_factory: None = ..., **kwargs: Any) -> _T_conn: ...
+@overload
 def connect(
-    dsn: Any | None = ..., connection_factory: Any | None = ..., cursor_factory: Any | None = ..., **kwargs
+    dsn: str | None = ..., *, connection_factory: Callable[..., _T_conn], cursor_factory: None = ..., **kwargs: Any
+) -> _T_conn: ...
+@overload
+def connect(
+    dsn: str | None = ...,
+    connection_factory: Callable[..., connection] | None = ...,
+    cursor_factory: Callable[..., cursor] | None = ...,
+    **kwargs: Any,
 ) -> connection: ...