add support for "async with cursor_context()"

aiopg/pool.py

@@ -256,15 +256,12 @@ def release(self, conn):
             fut = ensure_future(self._wakeup(), loop=self._loop)
         return fut
 
-    @asyncio.coroutine
-    def cursor(self, name=None, cursor_factory=None,
-               scrollable=None, withhold=False, *, timeout=None):
-        """XXX"""
-        conn = yield from self.acquire()
-        cur = yield from conn.cursor(name=name, cursor_factory=cursor_factory,
-                                     scrollable=scrollable, withhold=withhold,
-                                     timeout=timeout)
-        return _PoolCursorContextManager(self, conn, cur)
+    def cursor(self, name=None, cursor_factory=None, scrollable=None,
+               withhold=False, *, timeout=None):
+        cursor_kwargs = dict(name=name, cursor_factory=cursor_factory,
+                             scrollable=scrollable, withhold=withhold,
+                             timeout=timeout)
+        return _PoolCursorContextManager(self, cursor_kwargs)
 
     def __enter__(self):
         raise RuntimeError(

aiopg/utils.py

@@ -1,6 +1,6 @@
 import asyncio
 import sys
-
+import warnings
 
 PY_35 = sys.version_info >= (3, 5)
 PY_352 = sys.version_info >= (3, 5, 2)
@@ -125,7 +125,7 @@ class _PoolAcquireContextManager(_ContextManager):
     __slots__ = ('_coro', '_conn', '_pool')
 
     def __init__(self, coro, pool):
-        self._coro = coro
+        super().__init__(coro)
         self._conn = None
         self._pool = pool
 
@@ -163,6 +163,10 @@ def __init__(self, pool, conn):
         self._pool = pool
         self._conn = conn
 
+    @property
+    def conn(self):
+        return self._conn
+
     def __enter__(self):
         assert self._conn
         return self._conn
@@ -193,37 +197,97 @@ def __aexit__(self, exc_type, exc_val, exc_tb):
 class _PoolCursorContextManager:
     """Context manager.
 
-    This enables the following idiom for acquiring and releasing a
+    This enables the following idioms for acquiring and releasing a
     cursor around a block:
 
         with (yield from pool.cursor()) as cur:
             yield from cur.execute("SELECT 1")
 
+        async with pool.cursor() as cur:
+            yield from cur.execute("SELECT 1")
+
     while failing loudly when accidentally using:
 
         with pool:
             <block>
     """
 
-    __slots__ = ('_pool', '_conn', '_cur')
+    __slots__ = ('_pool', '_cursor_kwargs', '_cur')
 
-    def __init__(self, pool, conn, cur):
+    def __init__(self, pool, cursor_kwargs=None):
         self._pool = pool
-        self._conn = conn
-        self._cur = cur
+        self._cursor_kwargs = cursor_kwargs
+        self._cur = None
 
     def __enter__(self):
         return self._cur
 
     def __exit__(self, *args):
         try:
             self._cur.close()
-            self._pool.release(self._conn)
+            self._pool.__exit__(*args)
         finally:
-            self._pool = None
-            self._conn = None
             self._cur = None
 
+    @asyncio.coroutine
+    def _init_pool(self, with_aenter):
+        assert not self._cur
+
+        if with_aenter:
+            conn = None
+        else:
+            conn = yield from self._pool.acquire()
+
+        # self._pool now morphs into a _PoolConnectionContextManager
+        self._pool = _PoolConnectionContextManager(self._pool, conn)
+
+    @asyncio.coroutine
+    def _yield_await_init(self):
+        yield from self._init_pool(False)
+        self._cur = yield from self._pool.conn.cursor(**self._cursor_kwargs)
+        return self
+
+    def __iter__(self):
+        # This will get hit if you use "yield from pool.cursor()"
+        if PY_35:
+            warnings.warn("This usage is deprecated, use 'async with` syntax",
+                          DeprecationWarning)
+        return self._yield_await_init()
+
+    def __await__(self):
+        # This will get hit directly if you "await pool.cursor()"
+        # this is using a trick similar to the one here:
+        # https://magicstack.github.io/asyncpg/current/_modules/asyncpg/pool.html
+        # however since `self._init()` is an "asyncio.coroutine" we can't use
+        # just return self._init().__await__() as that returns a generator
+        # witout an "__await__" attribute, and we can't return a coroutine from
+        # here
+        if PY_35:
+            warnings.warn("This usage is deprecated, use 'async with` syntax",
+                          DeprecationWarning)
+        value = yield from self._yield_await_init()
+        return value
+
+    if PY_35:
+        @asyncio.coroutine
+        def __aenter__(self):
+            yield from self._init_pool(True)
+
+            # this will create the connection
+            yield from self._pool.__aenter__()
+            self._cur = yield from self._pool.conn.cursor(
+                **self._cursor_kwargs)
+            return self._cur
+
+        @asyncio.coroutine
+        def __aexit__(self, exc_type, exc_val, exc_tb):
+            try:
+                yield from self._cur.__aexit__(exc_type, exc_val, exc_tb)
+                self._cur = None
+            finally:
+                yield from self._pool.__aexit__(exc_type, exc_val, exc_tb)
+                self._pool = None
+
 
 if not PY_35:
     try:

tests/pep492/test_async_await.py

@@ -1,11 +1,9 @@
-import asyncio
 import pytest
 import aiopg
 import aiopg.sa
 from aiopg.sa import SAConnection
 
 
-@asyncio.coroutine
 async def test_cursor_await(make_connection):
     conn = await make_connection()
 
@@ -16,7 +14,6 @@ async def test_cursor_await(make_connection):
     cursor.close()
 
 
-@asyncio.coroutine
 async def test_connect_context_manager(loop, pg_params):
     async with aiopg.connect(loop=loop, **pg_params) as conn:
         cursor = await conn.cursor()
@@ -27,7 +24,26 @@ async def test_connect_context_manager(loop, pg_params):
     assert conn.closed
 
 
-@asyncio.coroutine
+async def test_pool_cursor_context_manager(loop, pg_params):
+    async with aiopg.create_pool(loop=loop, **pg_params) as pool:
+        async with pool.cursor() as cursor:
+            await cursor.execute('SELECT 42')
+            resp = await cursor.fetchone()
+            assert resp == (42, )
+        assert cursor.closed
+    assert pool.closed
+
+
+async def test_pool_cursor_await_context_manager(loop, pg_params):
+    async with aiopg.create_pool(loop=loop, **pg_params) as pool:
+        with (await pool.cursor()) as cursor:
+            await cursor.execute('SELECT 42')
+            resp = await cursor.fetchone()
+            assert resp == (42, )
+        assert cursor.closed
+    assert pool.closed
+
+
 async def test_connection_context_manager(make_connection):
     conn = await make_connection()
     assert not conn.closed
@@ -40,7 +56,6 @@ async def test_connection_context_manager(make_connection):
     assert conn.closed
 
 
-@asyncio.coroutine
 async def test_cursor_create_with_context_manager(make_connection):
     conn = await make_connection()
 
@@ -53,7 +68,6 @@ async def test_cursor_create_with_context_manager(make_connection):
     assert cursor.closed
 
 
-@asyncio.coroutine
 async def test_cursor_with_context_manager(make_connection):
     conn = await make_connection()
     cursor = await conn.cursor()
@@ -66,7 +80,6 @@ async def test_cursor_with_context_manager(make_connection):
     assert cursor.closed
 
 
-@asyncio.coroutine
 async def test_cursor_lightweight(make_connection):
     conn = await make_connection()
     cursor = await conn.cursor()
@@ -78,7 +91,6 @@ async def test_cursor_lightweight(make_connection):
     assert cursor.closed
 
 
-@asyncio.coroutine
 async def test_pool_context_manager(pg_params, loop):
     pool = await aiopg.create_pool(loop=loop, **pg_params)
 
@@ -93,7 +105,6 @@ async def test_pool_context_manager(pg_params, loop):
     assert pool.closed
 
 
-@asyncio.coroutine
 async def test_create_pool_context_manager(pg_params, loop):
     async with aiopg.create_pool(loop=loop, **pg_params) as pool:
         async with pool.acquire() as conn:
@@ -107,7 +118,6 @@ async def test_create_pool_context_manager(pg_params, loop):
     assert pool.closed
 
 
-@asyncio.coroutine
 async def test_cursor_aiter(make_connection):
     result = []
     conn = await make_connection()
@@ -122,7 +132,6 @@ async def test_cursor_aiter(make_connection):
     assert conn.closed
 
 
-@asyncio.coroutine
 async def test_engine_context_manager(pg_params, loop):
     engine = await aiopg.sa.create_engine(loop=loop, **pg_params)
     async with engine:
@@ -132,15 +141,13 @@ async def test_engine_context_manager(pg_params, loop):
     assert engine.closed
 
 
-@asyncio.coroutine
 async def test_create_engine_context_manager(pg_params, loop):
     async with aiopg.sa.create_engine(loop=loop, **pg_params) as engine:
         async with engine.acquire() as conn:
             assert isinstance(conn, SAConnection)
     assert engine.closed
 
 
-@asyncio.coroutine
 async def test_result_proxy_aiter(pg_params, loop):
     sql = 'SELECT generate_series(1, 5);'
     result = []
@@ -154,7 +161,6 @@ async def test_result_proxy_aiter(pg_params, loop):
     assert conn.closed
 
 
-@asyncio.coroutine
 async def test_transaction_context_manager(pg_params, loop):
     sql = 'SELECT generate_series(1, 5);'
     result = []
@@ -181,7 +187,6 @@ async def test_transaction_context_manager(pg_params, loop):
     assert conn.closed
 
 
-@asyncio.coroutine
 async def test_transaction_context_manager_error(pg_params, loop):
     async with aiopg.sa.create_engine(loop=loop, **pg_params) as engine:
         async with engine.acquire() as conn:
@@ -194,7 +199,6 @@ async def test_transaction_context_manager_error(pg_params, loop):
     assert conn.closed
 
 
-@asyncio.coroutine
 async def test_transaction_context_manager_commit_once(pg_params, loop):
     async with aiopg.sa.create_engine(loop=loop, **pg_params) as engine:
         async with engine.acquire() as conn:
@@ -214,7 +218,6 @@ async def test_transaction_context_manager_commit_once(pg_params, loop):
     assert conn.closed
 
 
-@asyncio.coroutine
 async def test_transaction_context_manager_nested_commit(pg_params, loop):
     sql = 'SELECT generate_series(1, 5);'
     result = []
@@ -244,7 +247,6 @@ async def test_transaction_context_manager_nested_commit(pg_params, loop):
     assert conn.closed
 
 
-@asyncio.coroutine
 async def test_sa_connection_execute(pg_params, loop):
     sql = 'SELECT generate_series(1, 5);'
     result = []