From: Denis Laxalde <denis.laxalde@dalibo.com>
Date: Wed, 5 Apr 2023 11:39:53 +0000 (+0200)
Subject: fix: use Self type for connection pool in __*enter__()
X-Git-Tag: pool-3.1.9~15
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=9c23dbd6f4b900649f14f079ad603868cbc32020;p=thirdparty%2Fpsycopg.git

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 87605755b..714c12180 100644
--- 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 1a1a85266..354cbd77c 100644
--- 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 27c5dcba2..a03b33d2f 100644
--- 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