Async context manager (#17)

* Implement async context manager support

* Fix tests

* Update CHANGES

Author: asvetlov

CHANGES.rst

@@ -9,6 +9,8 @@ CHANGES
 
 * Introduce `.timeout` property (#16)
 
+* Add methods for using as async context manager (#9)
+
 1.2.1 (2017-05-02)
 ------------------
 

README.rst

@@ -16,7 +16,7 @@ because ``timeout`` doesn't create a new task.
 The ``timeout(timeout, *, loop=None)`` call returns a context manager
 that cancels a block on *timeout* expiring::
 
-   with timeout(1.5):
+   async with timeout(1.5):
        await inner()
 
 1. If ``inner()`` is executed faster than in ``1.5`` seconds nothing
@@ -31,7 +31,7 @@ that cancels a block on *timeout* expiring::
 Context manager has ``.expired`` property for check if timeout happens
 exactly in context manager::
 
-   with timeout(1.5) as cm:
+   async with timeout(1.5) as cm:
        await inner()
    print(cm.expired)
 

async_timeout/__init__.py

@@ -30,6 +30,24 @@ def __init__(self, timeout, *, loop=None):
         self._cancel_handler = None
 
     def __enter__(self):
+        return self._do_enter()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._do_exit(exc_type)
+
+    @asyncio.coroutine
+    def __aenter__(self):
+        return self._do_enter()
+
+    @asyncio.coroutine
+    def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._do_exit(exc_type)
+
+    @property
+    def expired(self):
+        return self._cancelled
+
+    def _do_enter(self):
         if self._timeout is not None:
             self._task = current_task(self._loop)
             if self._task is None:
@@ -39,7 +57,7 @@ def __enter__(self):
                 self._timeout, self._cancel_task)
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def _do_exit(self, exc_type):
         if exc_type is asyncio.CancelledError and self._cancelled:
             self._cancel_handler = None
             self._task = None
@@ -53,9 +71,6 @@ def _cancel_task(self):
         self._task.cancel()
         self._cancelled = True
 
-    @property
-    def expired(self):
-        return self._cancelled
 
 
 def current_task(loop):

tests/conftest.py

@@ -0,0 +1,7 @@
+import sys
+
+
+def pytest_ignore_collect(path, config):
+    if 'py35' in str(path):
+        if sys.version_info < (3, 5, 0):
+            return True

tests/test_py35.py

@@ -0,0 +1,18 @@
+import asyncio
+
+import pytest
+
+from async_timeout import timeout
+
+
+async def test_async_timeout(loop):
+    with pytest.raises(asyncio.TimeoutError):
+        async with timeout(0.01, loop=loop) as cm:
+            await asyncio.sleep(10, loop=loop)
+    assert cm.expired
+
+
+async def test_async_no_timeout(loop):
+    async with timeout(1, loop=loop) as cm:
+        await asyncio.sleep(0, loop=loop)
+    assert not cm.expired