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 87605755b666988abb390f3061edfb76e33b4430..714c12180869dd1c6c1bf78238af63465780e046 100644 (file)
--- a/docs/news_pool.rst
+++ b/docs/news_pool.rst
@@ -7,6 +7,16 @@
 ``psycopg_pool`` release notes
 ==============================
 
+Future releases
+---------------
+
+psycopg_pool 3.1.9 (unreleased)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- Fix the return type annotation of `!NullConnectionPool.__enter__()`
+  (part of :ticket:`#540`).
+
+
 Current release
 ---------------
 

diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 1a1a85266c3edefd3a2a5937cce4c66a1a5f1389..354cbd77cd504c20d08ca8e5d3831df9b3cc3a5d 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 weakref import ref
 from contextlib import contextmanager
 
@@ -28,6 +28,8 @@ logger = logging.getLogger("psycopg.pool")
 
 
 class ConnectionPool(BasePool[Connection[Any]]):
+    _Self = TypeVar("_Self", bound="ConnectionPool")
+
     def __init__(
         self,
         conninfo: str = "",
@@ -383,7 +385,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 27c5dcba22b2dcc6e2e6cd9933863fd0d061a00a..a03b33d2ff6b9c8e042a5eed1f19107ca3248928 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
+from typing import Dict, List, Optional, Sequence, Type, TypeVar
 from weakref import ref
 from contextlib import asynccontextmanager
 
@@ -27,6 +27,8 @@ logger = logging.getLogger("psycopg.pool")
 
 
 class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
+    _Self = TypeVar("_Self", bound="AsyncConnectionPool")
+
     def __init__(
         self,
         conninfo: str = "",
@@ -327,7 +329,7 @@ class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
                 timeout,
             )
 
-    async def __aenter__(self) -> "AsyncConnectionPool":
+    async def __aenter__(self: _Self) -> _Self:
         await self.open()
         return self