Enable 'async for' for StreamReader, 'async with' for locks in 3.5.

2 files changed, 67 insertions, 55 deletions

diff --git a/asyncio/locks.py b/asyncio/locks.py
index 41a68c6..b2e516b 100644
--- a/asyncio/locks.py
+++ b/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/asyncio/streams.py b/asyncio/streams.py
index 64ff3d2..176c65e 100644
--- a/asyncio/streams.py
+++ b/asyncio/streams.py
@@ -6,6 +6,7 @@ __all__ = ['StreamReader', 'StreamWriter', 'StreamReaderProtocol',
            ]
 
 import socket
+import sys
 
 if hasattr(socket, 'AF_UNIX'):
     __all__.extend(['open_unix_connection', 'start_unix_server'])
@@ -19,6 +20,7 @@ from .log import logger
 
 
 _DEFAULT_LIMIT = 2**16
+_PY35 = sys.version_info >= (3, 5)
 
 
 class IncompleteReadError(EOFError):
@@ -485,3 +487,15 @@ class StreamReader:
             n -= len(block)
 
         return b''.join(blocks)
+
+    if _PY35:
+        @coroutine
+        def __aiter__(self):
+            return self
+
+        @coroutine
+        def __anext__(self):
+            val = yield from self.readline()
+            if val == b'':
+                raise StopAsyncIteration
+            return val