fix: use Self type for connection pool in __*enter__()

This makes inheritance, e.g. for *NullConnectionPool, work correctly:

    with psycopg_pool.NullConnectionPool() as p:
        pass
    reveal_type(p)
    # Revealed type is "psycopg_pool.null_pool.NullConnectionPool"

diff --git a/docs/news_pool.rst b/docs/news_pool.rst
index 0982e32632e0db00ce17d94f2d22e16c4e3bad47..1f2aeb8fd32b847e42c0ece22b2d581505236d87 100644 (file)
--- a/docs/news_pool.rst
+++ b/docs/news_pool.rst
@@ -17,6 +17,13 @@ psycopg_pool 3.2.0 (unreleased)
   (:ticket:`#520`).
 
 
+psycopg_pool 3.1.9 (unreleased)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- Fix the return type annotation of `!NullConnectionPool.__enter__()`
+  (:ticket:`#540`).
+
+
 Current release
 ---------------
 

diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 023f071646af4ad3aee672a3e556fd022fa398f3..47ecf084eb12b6447faf9d308331d50517877cd2 100644 (file)
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -11,7 +11,7 @@ from time import monotonic
 from queue import Queue, Empty
 from types import TracebackType
 from typing import Any, Callable, Dict, Iterator, List
-from typing import Optional, Sequence, Type
+from typing import Optional, Sequence, Type, TypeVar
 from typing_extensions import TypeAlias
 from weakref import ref
 from contextlib import contextmanager
@@ -31,6 +31,8 @@ ConnectFailedCB: TypeAlias = Callable[["ConnectionPool"], None]
 
 
 class ConnectionPool(BasePool[Connection[Any]]):
+    _Self = TypeVar("_Self", bound="ConnectionPool")
+
     def __init__(
         self,
         conninfo: str = "",
@@ -388,7 +390,7 @@ class ConnectionPool(BasePool[Connection[Any]]):
                         timeout,
                     )
 
-    def __enter__(self) -> "ConnectionPool":
+    def __enter__(self: _Self) -> _Self:
         self.open()
         return self
 

diff --git a/psycopg_pool/psycopg_pool/pool_async.py b/psycopg_pool/psycopg_pool/pool_async.py
index 84e90e46dd6799619a99855006f2d1f6b3f8c3b3..b57098e4e5baaf4bf6bc86aa72b82e4a43c36259 100644 (file)
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -10,7 +10,7 @@ from abc import ABC, abstractmethod
 from time import monotonic
 from types import TracebackType
 from typing import Any, AsyncIterator, Awaitable, Callable
-from typing import Dict, List, Optional, Sequence, Type, Union
+from typing import Dict, List, Optional, Sequence, Type, TypeVar, Union
 from typing_extensions import TypeAlias
 from weakref import ref
 from contextlib import asynccontextmanager
@@ -33,6 +33,8 @@ AsyncConnectFailedCB: TypeAlias = Union[
 
 
 class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
+    _Self = TypeVar("_Self", bound="AsyncConnectionPool")
+
     def __init__(
         self,
         conninfo: str = "",
@@ -333,7 +335,7 @@ class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
                 timeout,
             )
 
-    async def __aenter__(self) -> "AsyncConnectionPool":
+    async def __aenter__(self: _Self) -> _Self:
         await self.open()
         return self