0 is no-op timeout

aio-libs · Mar 13, 2017 · 3448389 · Kentzo · Sep 28, 2017 · 3448389
1 parent dfc3eec
commit 3448389
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 6 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -1,11 +1,13 @@
 CHANGES
 =======
 
-1.1.1 (2017-03-11)
+1.2.0 (2017-03-11)
 ------------------
 
 * Extra check on context manager exit
 
+* 0 is no-op timeout
+
 
 1.1.0 (2016-10-20)
 ------------------

diff --git a/async_timeout/__init__.py b/async_timeout/__init__.py
@@ -1,7 +1,7 @@
 import asyncio
 
 
-__version__ = '1.1.1'
+__version__ = '1.2.0'
 
 
 class timeout:
@@ -19,6 +19,8 @@ class timeout:
     loop - asyncio compatible event loop
     """
     def __init__(self, timeout, *, loop=None):
+        if timeout is not None and timeout == 0:
+            timeout = None
         self._timeout = timeout
         if loop is None:
             loop = asyncio.get_event_loop()
@@ -28,11 +30,11 @@ def __init__(self, timeout, *, loop=None):
         self._cancel_handler = None
 
     def __enter__(self):
-        self._task = asyncio.Task.current_task(loop=self._loop)
-        if self._task is None:
-            raise RuntimeError('Timeout context manager should be used '
-                               'inside a task')
         if self._timeout is not None:
+            self._task = asyncio.Task.current_task(loop=self._loop)
+            if self._task is None:
+                raise RuntimeError('Timeout context manager should be used '
+                                   'inside a task')
             self._cancel_handler = self._loop.call_later(
                 self._timeout, self._cancel_task)
         return self

diff --git a/tests/test_timeout.py b/tests/test_timeout.py
@@ -80,6 +80,21 @@ def long_running_task():
     assert 0.09 < dt < 0.13, dt
 
 
+@asyncio.coroutine
+def test_timeout_disable_zero(loop):
+    @asyncio.coroutine
+    def long_running_task():
+        yield from asyncio.sleep(0.1, loop=loop)
+        return 'done'
+
+    t0 = loop.time()
+    with timeout(0, loop=loop):
+        resp = yield from long_running_task()
+        assert resp == 'done'
+        dt = loop.time() - t0
+        assert 0.09 < dt < 0.13, dt
+
+
 @asyncio.coroutine
 def test_timeout_not_relevant_exception(loop):
     yield from asyncio.sleep(0, loop=loop)