Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Timeout context manager #611

Merged
merged 12 commits into from
Nov 25, 2015
44 changes: 44 additions & 0 deletions aiohttp/helpers.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
"""Various helper functions"""
import asyncio
import base64
import datetime
import functools
Expand Down Expand Up @@ -460,3 +461,46 @@ def requote_uri(uri):
# there may be unquoted '%'s in the URI. We need to make sure they're
# properly quoted so they do not cause issues elsewhere.
return quote(uri, safe=safe_without_percent)


class Timeout:
"""Timeout context manager.

Useful in cases when you want to apply timeout logic around block
of code or in cases when asyncio.wait_for is not suitable.

:param timeout: time out time in seconds
:param raise_error: if set, TimeoutError is raised in case of timeout
:param loop: asyncio compatible event loop
"""
def __init__(self, timeout, *, raise_error=False, loop=None):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

raise_error should be True by default (asyncio.wait_for() raises exception on timeout).
After thinking I doubt do we need the parameter at all.

self._timeout = timeout
if loop is None:
loop = asyncio.get_event_loop()
self._loop = loop
self._raise_error = raise_error
self._task = None
self._cancelled = False
self._cancel_handler = None

@asyncio.coroutine
def __aenter__(self):
self._task = asyncio.Task.current_task(loop=self._loop)
self._cancel_handler = self._loop.call_later(
self._timeout, self._cancel_task)
return self

@asyncio.coroutine
def __aexit__(self, exc_type, exc_val, exc_tb):
if self._cancelled:
if self._raise_error:
raise asyncio.TimeoutError
else:
# suppress
self._task = None
return True
else:
self._cancel_handler.cancel()

def _cancel_task(self):
self._cancelled = self._task.cancel()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add self._task = None

113 changes: 113 additions & 0 deletions tests/test_py35/test_timeout.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,113 @@
import asyncio
import time

import pytest
from aiohttp.helpers import Timeout


def test_timeout_without_error_raising(loop):
canceled_raised = False

async def long_running_task():
try:
await asyncio.sleep(10, loop=loop)
except asyncio.CancelledError:
nonlocal canceled_raised
canceled_raised = True
raise

async def run():
async with Timeout(0.01, loop=loop) as t:
await long_running_task()
assert t._loop is loop
assert canceled_raised, 'CancelledError was not raised'

loop.run_until_complete(run())


def test_timeout_raise_error(loop):
async def long_running_task():
await asyncio.sleep(10, loop=loop)

async def run():
with pytest.raises(asyncio.TimeoutError):
async with Timeout(0.01, raise_error=True, loop=loop):
await long_running_task()

loop.run_until_complete(run())


def test_timeout_finish_in_time(loop):
async def long_running_task():
await asyncio.sleep(0.01, loop=loop)
return 'done'

async def run():
async with Timeout(0.1, raise_error=True, loop=loop):
resp = await long_running_task()
assert resp == 'done'

loop.run_until_complete(run())


def test_timeout_gloabal_loop(loop):
asyncio.set_event_loop(loop)

async def run():
async with Timeout(0.1) as t:
await asyncio.sleep(0.01)
assert t._loop is loop

loop.run_until_complete(run())


def test_timeout_not_relevant_exception(loop):
async def run():
with pytest.raises(KeyError):
async with Timeout(0.1, loop=loop):
raise KeyError

loop.run_until_complete(run())


def test_timeout_blocking_loop(loop):
async def long_running_task():
time.sleep(0.1)
return 'done'

async def run():
async with Timeout(0.01, raise_error=True, loop=loop):
result = await long_running_task()
assert result == 'done'

loop.run_until_complete(run())


def test_for_race_conditions(loop):
async def run():
fut = asyncio.Future(loop=loop)
loop.call_later(0.1, fut.set_result('done'))
async with Timeout(0.2, raise_error=True, loop=loop):
resp = await fut
assert resp == 'done'

loop.run_until_complete(run())


def test_timeout_time(loop):
async def go():
foo_running = None

start = loop.time()
with pytest.raises(asyncio.TimeoutError):
async with Timeout(0.1, raise_error=True, loop=loop):
foo_running = True
try:
await asyncio.sleep(0.2, loop=loop)
finally:
foo_running = False

assert abs(0.1 - (loop.time() - start)) < 0.01
assert not foo_running

loop.run_until_complete(go())