Skip to content

Commit

Permalink
Allow to pass None as a timeout value to disable timeout logic
Browse files Browse the repository at this point in the history
In some cases there is needed to disable timeout logic but keep it in
convenient code structure without ugly if-branching.
  • Loading branch information
alexkey committed Mar 19, 2016
1 parent a1f39a9 commit 06e0e52
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 5 deletions.
12 changes: 7 additions & 5 deletions aiohttp/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -460,7 +460,7 @@ class Timeout:
... await r.text()
:param timeout: timeout value in seconds
:param timeout: timeout value in seconds or None to disable timeout logic
:param loop: asyncio compatible event loop
"""
def __init__(self, timeout, *, loop=None):
Expand All @@ -477,17 +477,19 @@ def __enter__(self):
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)
if self._timeout:
self._cancel_handler = self._loop.call_later(
self._timeout, self._cancel_task)
return self

def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is asyncio.CancelledError and self._cancelled:
self._cancel_handler = None
self._task = None
raise asyncio.TimeoutError
self._cancel_handler.cancel()
self._cancel_handler = None
if self._timeout:
self._cancel_handler.cancel()
self._cancel_handler = None
self._task = None

def _cancel_task(self):
Expand Down
13 changes: 13 additions & 0 deletions tests/test_timeout.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,19 @@ def run():
loop.run_until_complete(run())


@pytest.mark.parametrize("delay", [None, 0])
@pytest.mark.run_loop
def test_timeout_disable(loop, delay):
@asyncio.coroutine
def long_running_task():
yield from asyncio.sleep(0.1, loop=loop)
return 'done'

with Timeout(delay, loop=loop):
resp = yield from long_running_task()
assert resp == 'done'


@pytest.mark.run_loop
def test_timeout_not_relevant_exception(loop):
yield from asyncio.sleep(0, loop=loop)
Expand Down

0 comments on commit 06e0e52

Please sign in to comment.