Issue 24178: support 'async with' for asyncio locks.

3 files changed, 124 insertions, 55 deletions

diff --git a/Lib/asyncio/locks.py b/Lib/asyncio/locks.py
index 41a68c6..b2e516b 100644
--- a/Lib/asyncio/locks.py
+++ b/Lib/asyncio/locks.py
@@ -3,12 +3,16 @@
 __all__ = ['Lock', 'Event', 'Condition', 'Semaphore', 'BoundedSemaphore']
 
 import collections
+import sys
 
 from . import events
 from . import futures
 from .coroutines import coroutine
 
 
+_PY35 = sys.version_info >= (3, 5)
+
+
 class _ContextManager:
     """Context manager.
 
@@ -39,7 +43,53 @@ class _ContextManager:
             self._lock = None  # Crudely prevent reuse.
 
 
-class Lock:
+class _ContextManagerMixin:
+    def __enter__(self):
+        raise RuntimeError(
+            '"yield from" should be used as context manager expression')
+
+    def __exit__(self, *args):
+        # This must exist because __enter__ exists, even though that
+        # always raises; that's how the with-statement works.
+        pass
+
+    @coroutine
+    def __iter__(self):
+        # This is not a coroutine.  It is meant to enable the idiom:
+        #
+        #     with (yield from lock):
+        #         <block>
+        #
+        # as an alternative to:
+        #
+        #     yield from lock.acquire()
+        #     try:
+        #         <block>
+        #     finally:
+        #         lock.release()
+        yield from self.acquire()
+        return _ContextManager(self)
+
+    if _PY35:
+
+        def __await__(self):
+            # To make "with await lock" work.
+            yield from self.acquire()
+            return _ContextManager(self)
+
+        @coroutine
+        def __aenter__(self):
+            yield from self.acquire()
+            # We have no use for the "as ..."  clause in the with
+            # statement for locks.
+            return None
+
+        @coroutine
+        def __aexit__(self, exc_type, exc, tb):
+            self.release()
+
+
+class Lock(_ContextManagerMixin):
     """Primitive lock objects.
 
     A primitive lock is a synchronization primitive that is not owned
@@ -153,32 +203,6 @@ class Lock:
         else:
             raise RuntimeError('Lock is not acquired.')
 
-    def __enter__(self):
-        raise RuntimeError(
-            '"yield from" should be used as context manager expression')
-
-    def __exit__(self, *args):
-        # This must exist because __enter__ exists, even though that
-        # always raises; that's how the with-statement works.
-        pass
-
-    @coroutine
-    def __iter__(self):
-        # This is not a coroutine.  It is meant to enable the idiom:
-        #
-        #     with (yield from lock):
-        #         <block>
-        #
-        # as an alternative to:
-        #
-        #     yield from lock.acquire()
-        #     try:
-        #         <block>
-        #     finally:
-        #         lock.release()
-        yield from self.acquire()
-        return _ContextManager(self)
-
 
 class Event:
     """Asynchronous equivalent to threading.Event.
@@ -246,7 +270,7 @@ class Event:
             self._waiters.remove(fut)
 
 
-class Condition:
+class Condition(_ContextManagerMixin):
     """Asynchronous equivalent to threading.Condition.
 
     This class implements condition variable objects. A condition variable
@@ -356,21 +380,8 @@ class Condition:
         """
         self.notify(len(self._waiters))
 
-    def __enter__(self):
-        raise RuntimeError(
-            '"yield from" should be used as context manager expression')
-
-    def __exit__(self, *args):
-        pass
 
-    @coroutine
-    def __iter__(self):
-        # See comment in Lock.__iter__().
-        yield from self.acquire()
-        return _ContextManager(self)
-
-
-class Semaphore:
+class Semaphore(_ContextManagerMixin):
     """A Semaphore implementation.
 
     A semaphore manages an internal counter which is decremented by each
@@ -441,19 +452,6 @@ class Semaphore:
                 waiter.set_result(True)
                 break
 
-    def __enter__(self):
-        raise RuntimeError(
-            '"yield from" should be used as context manager expression')
-
-    def __exit__(self, *args):
-        pass
-
-    @coroutine
-    def __iter__(self):
-        # See comment in Lock.__iter__().
-        yield from self.acquire()
-        return _ContextManager(self)
-
 
 class BoundedSemaphore(Semaphore):
     """A bounded semaphore implementation.
diff --git a/Lib/test/test_asyncio/test_pep492.py b/Lib/test/test_asyncio/test_pep492.py
new file mode 100644
index 0000000..c9a2f96
--- /dev/null
+++ b/Lib/test/test_asyncio/test_pep492.py
@@ -0,0 +1,68 @@
+"""Tests support for new syntax introduced by PEP 492."""
+
+import unittest
+from unittest import mock
+
+import asyncio
+from asyncio import test_utils
+
+
+class BaseTest(test_utils.TestCase):
+
+    def setUp(self):
+        self.loop = asyncio.BaseEventLoop()
+        self.loop._process_events = mock.Mock()
+        self.loop._selector = mock.Mock()
+        self.loop._selector.select.return_value = ()
+        self.set_event_loop(self.loop)
+
+
+class LockTests(BaseTest):
+
+    def test_context_manager_async_with(self):
+        primitives = [
+            asyncio.Lock(loop=self.loop),
+            asyncio.Condition(loop=self.loop),
+            asyncio.Semaphore(loop=self.loop),
+            asyncio.BoundedSemaphore(loop=self.loop),
+        ]
+
+        async def test(lock):
+            await asyncio.sleep(0.01, loop=self.loop)
+            self.assertFalse(lock.locked())
+            async with lock as _lock:
+                self.assertIs(_lock, None)
+                self.assertTrue(lock.locked())
+                await asyncio.sleep(0.01, loop=self.loop)
+                self.assertTrue(lock.locked())
+            self.assertFalse(lock.locked())
+
+        for primitive in primitives:
+            self.loop.run_until_complete(test(primitive))
+            self.assertFalse(primitive.locked())
+
+    def test_context_manager_with_await(self):
+        primitives = [
+            asyncio.Lock(loop=self.loop),
+            asyncio.Condition(loop=self.loop),
+            asyncio.Semaphore(loop=self.loop),
+            asyncio.BoundedSemaphore(loop=self.loop),
+        ]
+
+        async def test(lock):
+            await asyncio.sleep(0.01, loop=self.loop)
+            self.assertFalse(lock.locked())
+            with await lock as _lock:
+                self.assertIs(_lock, None)
+                self.assertTrue(lock.locked())
+                await asyncio.sleep(0.01, loop=self.loop)
+                self.assertTrue(lock.locked())
+            self.assertFalse(lock.locked())
+
+        for primitive in primitives:
+            self.loop.run_until_complete(test(primitive))
+            self.assertFalse(primitive.locked())
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/Misc/NEWS b/Misc/NEWS
index cba5432..b7cd1c7 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -116,6 +116,9 @@ Library
 
 - asyncio: async() function is deprecated in favour of ensure_future().
 
+- Issue 24178: asyncio.Lock, Condition, Semaphore, and BoundedSemaphore
+  support new 'async with' syntax.  Contributed by Yury Selivanov.
+
 Tests
 -----