diff --git a/README.md b/README.md index f93b0de..3689113 100644 --- a/README.md +++ b/README.md @@ -186,6 +186,25 @@ ic| 3: 3 code with `ic()` breaks. +### Timing Asynchronous Functions + +`@ic.timer` measures a coroutine function from when it is awaited until it +returns, raises, or is cancelled. The decorated function remains a coroutine +function and preserves its return value and exceptions. + +```python +import asyncio +from icecream import ic + +@ic.timer +async def pause(): + await asyncio.sleep(0.1) + +asyncio.run(pause()) +# ic| pause took 100.00ms (approximately) +``` + + ### Import Tricks To make `ic()` available in every file without needing to be imported in diff --git a/icecream/icecream.py b/icecream/icecream.py index 45a39ca..7a07d70 100644 --- a/icecream/icecream.py +++ b/icecream/icecream.py @@ -574,6 +574,18 @@ def _output(self, duration: float, label: str = "") -> None: self._ic.outputFunction(msg) def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]: + if inspect.iscoroutinefunction(func): + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + start_time = time.perf_counter() + try: + return await func(*args, **kwargs) + finally: + duration = time.perf_counter() - start_time + self._output(duration, func.__name__) + + return async_wrapper + @functools.wraps(func) def wrapper(*args: Any, **kwargs: Any) -> Any: start_time: float = time.perf_counter() diff --git a/tests/test_icecream.py b/tests/test_icecream.py index 342c56d..1d80fae 100644 --- a/tests/test_icecream.py +++ b/tests/test_icecream.py @@ -9,11 +9,14 @@ # # License: MIT # +import asyncio +import inspect import re import sys import time import unittest import warnings +from unittest.mock import patch from io import StringIO from contextlib import contextmanager @@ -994,3 +997,84 @@ def test_timer_context_manager_resets_enter_time_after_exit(self): with self.assertRaises(RuntimeError) as ctx: timer.__exit__(None, None, None) self.assertIn("__enter__", str(ctx.exception)) + + +class TestAsyncTimer(unittest.IsolatedAsyncioTestCase): + def test_preserves_coroutine_function(self): + @ic.timer + async def original(): + """An asynchronous function.""" + + self.assertTrue(inspect.iscoroutinefunction(original)) + self.assertEqual(original.__name__, 'original') + self.assertEqual(original.__doc__, 'An asynchronous function.') + + async def test_measures_awaited_function(self): + output = [] + debugger = icecream.IceCreamDebugger(outputFunction=output.append) + + @debugger.timer + async def add(x, y=0): + await asyncio.sleep(0) + self.assertEqual(output, []) + return x + y + + with patch('icecream.icecream.time.perf_counter', side_effect=[10, 12]): + coroutine = add(2, y=3) + try: + self.assertEqual(output, []) + finally: + result = await coroutine + + self.assertEqual(result, 5) + self.assertEqual(output, ['ic| add took 2.00s']) + + async def test_reports_time_and_preserves_exception(self): + output = [] + debugger = icecream.IceCreamDebugger(outputFunction=output.append) + error = ValueError('asynchronous failure') + + @debugger.timer + async def fail(): + await asyncio.sleep(0) + raise error + + with patch('icecream.icecream.time.perf_counter', side_effect=[10, 12]): + with self.assertRaises(ValueError) as caught: + await fail() + + self.assertIs(caught.exception, error) + self.assertEqual(output, ['ic| fail took 2.00s']) + + async def test_disabled_timer(self): + output = [] + debugger = icecream.IceCreamDebugger(outputFunction=output.append) + debugger.disable() + + @debugger.timer + async def original(): + await asyncio.sleep(0) + return 42 + + self.assertEqual(await original(), 42) + self.assertEqual(output, []) + + async def test_reports_time_when_cancelled(self): + output = [] + debugger = icecream.IceCreamDebugger(outputFunction=output.append) + started = asyncio.Event() + + @debugger.timer + async def wait(): + started.set() + await asyncio.Event().wait() + + with patch('icecream.icecream.time.perf_counter', side_effect=[10, 12]): + task = asyncio.create_task(wait()) + await started.wait() + self.assertEqual(output, []) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + self.assertEqual(output, ['ic| wait took 2.00s'])