diff --git a/stubs/psycopg2/psycopg2/pool.pyi b/stubs/psycopg2/psycopg2/pool.pyi index c2a4dc47ab5a..a883d74f0ed8 100644 --- a/stubs/psycopg2/psycopg2/pool.pyi +++ b/stubs/psycopg2/psycopg2/pool.pyi @@ -2,6 +2,7 @@ from _typeshed import ConvertibleToInt from collections.abc import Hashable import psycopg2 +from psycopg2.extensions import connection class PoolError(psycopg2.Error): ... @@ -12,8 +13,8 @@ class AbstractConnectionPool: def __init__(self, minconn: ConvertibleToInt, maxconn: ConvertibleToInt, *args, **kwargs) -> None: ... # getconn, putconn and closeall are officially documented as methods of the # abstract base class, but in reality, they only exist on the children classes - def getconn(self, key: Hashable | None = None): ... - def putconn(self, conn, key: Hashable | None = None, close: bool = False) -> None: ... + def getconn(self, key: Hashable | None = None) -> connection: ... + def putconn(self, conn: connection, key: Hashable | None = None, close: bool = False) -> None: ... def closeall(self) -> None: ... class SimpleConnectionPool(AbstractConnectionPool): ... @@ -21,4 +22,4 @@ class SimpleConnectionPool(AbstractConnectionPool): ... class ThreadedConnectionPool(AbstractConnectionPool): # This subclass has a default value for conn which doesn't exist # in the SimpleConnectionPool class, nor in the documentation - def putconn(self, conn=None, key: Hashable | None = None, close: bool = False) -> None: ... + def putconn(self, conn: connection | None = None, key: Hashable | None = None, close: bool = False) -> None: ...