diff --git a/CHANGES/141.feature b/CHANGES/141.feature new file mode 100644 index 0000000..8c77334 --- /dev/null +++ b/CHANGES/141.feature @@ -0,0 +1 @@ +`timeout` can now also be used as function decorator. \ No newline at end of file diff --git a/README.rst b/README.rst index f1174d4..a6202be 100644 --- a/README.rst +++ b/README.rst @@ -37,6 +37,13 @@ that cancels a block on *timeout* expiring:: *timeout* parameter could be ``None`` for skipping timeout functionality. +Another possibility is to use ``timeout`` as function decorator:: + + @timeout(1.5) + async def possibly_long_running(): + pass + + Alternatively, ``timeout_at(when)`` can be used for scheduling at the absolute time:: diff --git a/async_timeout/__init__.py b/async_timeout/__init__.py index e86b319..51c5198 100644 --- a/async_timeout/__init__.py +++ b/async_timeout/__init__.py @@ -128,6 +128,13 @@ async def __aexit__( self._do_exit(exc_type) return None + def __call__(self, f): + async def inner(*args, **kargs): + async with self: + return await f(*args, **kargs) + + return inner + @property def expired(self) -> bool: """Is timeout expired during execution?""" diff --git a/tests/test_timeout.py b/tests/test_timeout.py index 28ad4be..41a0603 100644 --- a/tests/test_timeout.py +++ b/tests/test_timeout.py @@ -26,6 +26,26 @@ async def long_running_task(): assert canceled_raised, "CancelledError was not raised" +@pytest.mark.asyncio +async def test_timeout_as_decorator(): + canceled_raised = False + + @timeout(0.01) + async def long_running_task(arg, *, karg): + assert arg == 1 + assert karg == 2 + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + nonlocal canceled_raised + canceled_raised = True + raise + + with pytest.raises(asyncio.TimeoutError): + await long_running_task(1, karg=2) + assert canceled_raised, "CancelledError was not raised" + + @pytest.mark.asyncio async def test_timeout_finish_in_time(): async def long_running_task():