diff --git a/_python_utils_tests/clock.py b/_python_utils_tests/clock.py new file mode 100644 index 0000000..d0c5cbe --- /dev/null +++ b/_python_utils_tests/clock.py @@ -0,0 +1,56 @@ +"""A fake clock for the tests and doctests of ``python_utils.time``.""" + +import typing + +import pytest + +import python_utils.time + +#: The doctest that sleeps on the clock, by the name pytest gives it. +TIMEOUT_GENERATOR_DOCTEST: str = 'python_utils.time.timeout_generator' + + +class FakeClock: + """ + A clock that only moves when something sleeps on it. + + ``time.sleep`` promises to sleep at least as long as requested. A busy + machine sleeps tens of milliseconds longer, and that changes how many + items ``timeout_generator`` yields before its timeout. On this clock the + number of items depends on the arguments alone. + + Attributes: + now (float): The current time in seconds. + sleeps (list[float]): Every requested sleep, in order. + """ + + def __init__(self) -> None: + """Start at zero without any recorded sleeps.""" + self.now: float = 0.0 + self.sleeps: list[float] = [] + + def perf_counter(self) -> float: + """Return the current time, like ``time.perf_counter``.""" + return self.now + + def sleep(self, seconds: float) -> None: + """Record the sleep and move the clock forward, without waiting.""" + self.sleeps.append(seconds) + self.now += seconds + + +@pytest.fixture +def fake_clock(monkeypatch: pytest.MonkeyPatch) -> FakeClock: + """Replace the ``time`` module inside ``python_utils.time``.""" + clock: FakeClock = FakeClock() + monkeypatch.setattr(python_utils.time, 'time', clock) + return clock + + +@pytest.fixture(autouse=True) +def fake_clock_in_doctest(request: pytest.FixtureRequest) -> None: + """Run the ``timeout_generator`` doctest on the fake clock.""" + # pytest leaves `FixtureRequest.node` without a type. + node: pytest.Item = typing.cast(pytest.Item, request.node) + if node.name == TIMEOUT_GENERATOR_DOCTEST: + request.getfixturevalue('fake_clock') diff --git a/_python_utils_tests/test_aio.py b/_python_utils_tests/test_aio.py index e7532a8..f9461fe 100644 --- a/_python_utils_tests/test_aio.py +++ b/_python_utils_tests/test_aio.py @@ -25,6 +25,43 @@ async def mock_sleep(delay: float) -> None: assert sum(sleeps) == 4 +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'start,step,stop,expected', + [ + (10, -2, 0, [10, 8, 6, 4, 2]), + (0, -1, -3, [0, -1, -2]), + (1.5, -0.5, 0, [1.5, 1.0, 0.5]), + # Nothing to count when `start` is already past `stop`, in either + # direction. This is what `range` does as well. + (0, -1, 5, []), + (5, 1, 3, []), + (0, 2, 5, [0, 2, 4]), + ], +) +async def test_acount_stop_follows_step_direction( + start: float, step: float, stop: float, expected: types.List[float] +) -> None: + """Count down to a lower ``stop`` when the step is negative.""" + limit: int = len(expected) + 5 + items: types.List[float] = [] + item: float + # `acount` is annotated as a plain iterator, closing takes a generator. + counter: types.AsyncGenerator[float, None] = types.cast( + types.AsyncGenerator[float, None], + aio.acount(start=start, step=step, stop=stop), + ) + async for item in counter: + items.append(item) + # A counter that misses its `stop` never ends by itself. + if len(items) == limit: + break + + await counter.aclose() + + assert items == expected + + @pytest.mark.asyncio async def test_acontainer() -> None: """Collect an async iterable into the requested container.""" diff --git a/_python_utils_tests/test_decorators.py b/_python_utils_tests/test_decorators.py index 0f8e070..571f50f 100644 --- a/_python_utils_tests/test_decorators.py +++ b/_python_utils_tests/test_decorators.py @@ -1,5 +1,7 @@ """Tests for the decorators in ``python_utils.decorators``.""" +import inspect +import logging import typing from unittest import mock @@ -46,6 +48,50 @@ def test_sample_not_called(random: mock.MagicMock) -> None: assert demo_function.call_count == 0 +def test_sample_leaves_root_logger_alone( + random: mock.MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + """Skip a call without configuring the root logger.""" + root: logging.Logger = logging.getLogger() + handlers: list[logging.Handler] = [] + # A root logger without handlers is what an application starts with. + # `logging.debug()` installs a handler on it, a module logger does not. + monkeypatch.setattr(root, 'handlers', handlers) + random.return_value = 1.0 + + decorators.sample(0.5)(mock.MagicMock())() + + assert handlers == [] + + +def test_sample_logs_on_module_logger( + random: mock.MagicMock, caplog: pytest.LogCaptureFixture +) -> None: + """Report a skipped call on the logger of the module.""" + random.return_value = 1.0 + + with caplog.at_level(logging.DEBUG): + decorators.sample(0.5)(mock.MagicMock())() + + names: list[str] = [record.name for record in caplog.records] + assert names == ['python_utils.decorators'] + assert 'Skipped execution' in caplog.records[0].getMessage() + + +def test_listify_keeps_metadata() -> None: + """Keep the name, docstring and signature of the decorated function.""" + + @decorators.listify(collection=list) + def numbers(count: int = 3) -> typing.Iterator[int]: + """Yield ``count`` numbers.""" + yield from range(count) + + assert numbers() == [0, 1, 2] + assert numbers.__name__ == 'numbers' + assert numbers.__doc__ == 'Yield ``count`` numbers.' + assert list(inspect.signature(numbers).parameters) == ['count'] + + class SomeClass: """A sample class with classmethods for wrapping tests.""" @@ -80,3 +126,53 @@ def test_wraps_annotated_classmethod() -> None: )(some_class.some_annotated_classmethod) wrapped_method(123) some_class.some_annotated_classmethod.assert_called_with(123) + + +def test_wraps_classmethod_leaves_wrapped_annotations_alone() -> None: + """Drop ``self`` for the wrapper without touching the wrapped method.""" + + def wrapped(self: SomeClass, arg: int) -> int: + """Return the argument unchanged, as a regular method would.""" + return arg + + def wrapper(cls: type[SomeClass], arg: int) -> int: + """Return the argument unchanged, as a classmethod would.""" + return arg + + result: typing.Callable[..., int] = decorators.wraps_classmethod(wrapped)( + wrapper + ) + + assert wrapped.__annotations__ == { + 'self': SomeClass, + 'arg': int, + 'return': int, + } + assert result.__annotations__ == {'arg': int, 'return': int} + assert result.__annotations__ is not wrapped.__annotations__ + + +def test_wraps_classmethod_keeps_wrapper_annotations() -> None: + """Keep the wrapper's annotations if the wrapped method has none.""" + + def wrapped(self: SomeClass, arg: int) -> int: + """Return the argument unchanged, as a regular method would.""" + return arg + + def wrapper(cls: type[SomeClass], arg: int) -> int: + """Return the argument unchanged, as a classmethod would.""" + return arg + + # To the interpreter this is a method that was written without + # annotations. The type checkers still get to see them. + wrapped.__annotations__ = {} + + result: typing.Callable[..., int] = decorators.wraps_classmethod(wrapped)( + wrapper + ) + + assert result.__annotations__ == { + 'cls': type[SomeClass], + 'arg': int, + 'return': int, + } diff --git a/_python_utils_tests/test_generators.py b/_python_utils_tests/test_generators.py index c1b44d9..501db0f 100644 --- a/_python_utils_tests/test_generators.py +++ b/_python_utils_tests/test_generators.py @@ -1,6 +1,8 @@ """Tests for the batching helpers in ``python_utils.generators``.""" import asyncio +import gc +import sys from types import SimpleNamespace import pytest @@ -87,6 +89,217 @@ async def generator() -> types.AsyncIterator[int]: await batcher.__anext__() +async def blocked_generator( + started: asyncio.Event, closed: types.List[str] +) -> types.AsyncIterator[int]: + """Yield one item and then wait for an event that never comes.""" + try: + yield 0 + started.set() + await asyncio.Event().wait() + finally: + closed.append('closed') + + +@pytest.mark.asyncio +async def test_abatcher_cancels_pending_item_on_close() -> None: + """Cancel the pending source item when the consumer stops early.""" + closed: types.List[str] = [] + before: types.Set[asyncio.Task[types.Any]] = asyncio.all_tasks() + + # The source never yields a second item, so the interval always ends + # while that item is still pending. Sleep accuracy does not matter. + batcher: types.AsyncGenerator[types.List[int], None] = ( + python_utils.abatcher( + blocked_generator(asyncio.Event(), closed), interval=0.01 + ) + ) + first: types.List[int] = await batcher.__anext__() + await batcher.aclose() + + assert first == [0] + assert asyncio.all_tasks() == before + assert closed == ['closed'] + + +@pytest.mark.asyncio +async def test_abatcher_cancels_pending_item_on_cancellation() -> None: + """Cancel the pending source item when the consumer is cancelled.""" + started: asyncio.Event = asyncio.Event() + closed: types.List[str] = [] + before: types.Set[asyncio.Task[types.Any]] = asyncio.all_tasks() + + batcher: types.AsyncGenerator[types.List[int], None] = ( + python_utils.abatcher(blocked_generator(started, closed), batch_size=2) + ) + consumer: asyncio.Task[types.List[int]] = asyncio.create_task( + batcher.__anext__() + ) + await started.wait() + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + + assert asyncio.all_tasks() == before + assert closed == ['closed'] + + +@pytest.mark.asyncio +async def test_abatcher_reports_source_error_during_cancellation() -> None: + """Report an error that the source raises while it is cancelled.""" + reported: types.List[types.Dict[str, types.Any]] = [] + loop: asyncio.AbstractEventLoop = asyncio.get_running_loop() + loop.set_exception_handler(lambda _, context: reported.append(context)) + + async def generator() -> types.AsyncGenerator[int, None]: + """Yield one item, then fail while the next one is cancelled.""" + yield 0 + try: + await asyncio.Event().wait() + finally: + raise RuntimeError('closing failed') + + batcher: types.AsyncGenerator[types.List[int], None] = ( + python_utils.abatcher(generator(), interval=0.01) + ) + try: + first: types.List[int] = await batcher.__anext__() + await batcher.aclose() + finally: + loop.set_exception_handler(None) + + assert first == [0] + assert [str(context['exception']) for context in reported] == [ + 'closing failed' + ] + + +@pytest.mark.asyncio +async def test_abatcher_accepts_source_that_ends_on_cancellation() -> None: + """Report nothing for a source that just stops when it is cancelled.""" + reported: types.List[types.Dict[str, types.Any]] = [] + loop: asyncio.AbstractEventLoop = asyncio.get_running_loop() + loop.set_exception_handler(lambda _, context: reported.append(context)) + + async def generator() -> types.AsyncGenerator[int, None]: + """Yield one item, then end as soon as the wait is cancelled.""" + yield 0 + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + return + + batcher: types.AsyncGenerator[types.List[int], None] = ( + python_utils.abatcher(generator(), interval=0.01) + ) + try: + first: types.List[int] = await batcher.__anext__() + await batcher.aclose() + finally: + loop.set_exception_handler(None) + + assert first == [0] + assert reported == [] + + +def test_abatcher_collected_after_its_loop_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Clean up without an event loop when nothing is on its way.""" + unraisable: types.List[types.Any] = [] + monkeypatch.setattr(sys, 'unraisablehook', unraisable.append) + + async def generator() -> types.AsyncGenerator[int, None]: + """Yield four items without waiting.""" + i: int + for i in range(4): + yield i + + loop: asyncio.AbstractEventLoop = asyncio.new_event_loop() + batcher: types.AsyncGenerator[types.List[int], None] = ( + python_utils.abatcher(generator(), batch_size=2) + ) + try: + first: types.List[int] = loop.run_until_complete(batcher.__anext__()) + finally: + loop.close() + + del batcher + gc.collect() + + assert first == [0, 1] + assert unraisable == [] + + +@pytest.mark.asyncio +async def test_abatcher_size_flush_restarts_interval( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Restart the interval after a batch that was flushed by its size.""" + now: float = 0.0 + monkeypatch.setattr( + python_utils.generators, + 'time', + SimpleNamespace(perf_counter=lambda: now), + ) + + async def generator() -> types.AsyncIterator[int]: + """Let one item arrive every 4 seconds on the fake clock.""" + nonlocal now + item: int + for item in range(12): + now = 4.0 * (item + 1) + yield item + + batches: types.List[types.List[int]] = [ + batch + async for batch in python_utils.abatcher( + generator(), batch_size=3, interval=10 + ) + ] + + # Three items take 12 seconds, which is past the interval of 10. The size + # flush takes them and the next item has to start a fresh interval. It + # may not leave on its own because the old interval ran out. + assert batches == [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]] + + +class FutureIterator: + """An async iterator that hands out futures instead of coroutines.""" + + def __init__(self, stop: int) -> None: + """Count from zero up to ``stop``, which is excluded.""" + self.current: int = 0 + self.stop: int = stop + + def __aiter__(self) -> 'FutureIterator': + """Return the iterator itself.""" + return self + + def __anext__(self) -> 'asyncio.Future[int]': + """Return a finished future with the next number.""" + future: asyncio.Future[int] = ( + asyncio.get_running_loop().create_future() + ) + if self.current < self.stop: + future.set_result(self.current) + self.current += 1 + else: + future.set_exception(StopAsyncIteration()) + + return future + + +@pytest.mark.asyncio +async def test_abatcher_with_future_returning_iterator() -> None: + """Batch an iterator whose ``__anext__`` is not a coroutine.""" + batches: types.List[types.List[int]] = [ + batch async for batch in python_utils.abatcher(FutureIterator(5), 2) + ] + + assert batches == [[0, 1], [2, 3], [4]] + + def test_batcher() -> None: """Split an iterable into fixed-size batches.""" batch = [] diff --git a/_python_utils_tests/test_lazy_imports.py b/_python_utils_tests/test_lazy_imports.py index da8143e..54ec64c 100644 --- a/_python_utils_tests/test_lazy_imports.py +++ b/_python_utils_tests/test_lazy_imports.py @@ -11,6 +11,7 @@ import pytest import python_utils +from _python_utils_tests import clock def _run_clean(code: str) -> subprocess.CompletedProcess[str]: @@ -94,10 +95,13 @@ def test_star_import_resolves_all_names() -> None: @pytest.mark.asyncio -async def test_aio_timeout_generator_default_iterable() -> None: +async def test_aio_timeout_generator_default_iterable( + fake_clock: clock.FakeClock, +) -> None: """Default the iterable to ``aio.acount`` when omitted.""" # With no iterable the generator defaults to ``aio.acount`` -- exercising # the lazy ``aio``/``asyncio`` import and the None-resolution branch. + # The fake clock stands still, so the timeout cannot end the loop early. count = 0 generator: collections.abc.AsyncGenerator[object, None] = ( python_utils.aio_timeout_generator(timeout=0.05, interval=0.0) diff --git a/_python_utils_tests/test_time.py b/_python_utils_tests/test_time.py index da26c9e..b30d7c5 100644 --- a/_python_utils_tests/test_time.py +++ b/_python_utils_tests/test_time.py @@ -3,12 +3,18 @@ import asyncio import datetime import itertools +import time import pytest import python_utils +from _python_utils_tests import clock from python_utils import types +#: Far longer than every timeout in this module, so a generator that sleeps +#: this long is always interrupted first. +STALL: float = 10.0 + @pytest.mark.parametrize( 'timeout,interval,interval_multiplier,maximum_interval,iterable,result', @@ -28,6 +34,8 @@ ) @pytest.mark.asyncio async def test_aio_timeout_generator( + fake_clock: clock.FakeClock, + monkeypatch: pytest.MonkeyPatch, timeout: float, interval: float, interval_multiplier: float, @@ -36,6 +44,13 @@ async def test_aio_timeout_generator( result: int, ) -> None: """Stop the async generator near the configured timeout.""" + + async def sleep(delay: float) -> None: + """Let the fake clock pass the delay without waiting for it.""" + fake_clock.sleep(delay) + + monkeypatch.setattr(asyncio, 'sleep', sleep) + i = None async for i in python_utils.aio_timeout_generator( timeout, interval, iterable, maximum_interval=maximum_interval @@ -46,12 +61,13 @@ async def test_aio_timeout_generator( @pytest.mark.parametrize( - 'timeout,interval,interval_multiplier,maximum_interval,iterable,result', + 'timeout,interval,interval_multiplier,maximum_interval,iterable,result,' + 'sleeps', [ - (0.1, 0.06, 0.5, 0.1, 'abc', 'c'), - (0.1, 0.07, 0.5, 0.1, itertools.count, 2), - (0.1, 0.07, 0.5, 0.1, itertools.count(), 2), - (0.1, 0.06, 1.0, None, 'abc', 'c'), + (0.1, 0.06, 0.5, 0.1, 'abc', 'c', [0.06, 0.03, 0.015]), + (0.1, 0.07, 0.5, 0.1, itertools.count, 2, [0.07, 0.035]), + (0.1, 0.07, 0.5, 0.1, itertools.count(), 2, [0.07, 0.035]), + (0.1, 0.06, 1.0, None, 'abc', 'c', [0.06, 0.06]), ( datetime.timedelta(seconds=0.1), datetime.timedelta(seconds=0.06), @@ -59,10 +75,12 @@ async def test_aio_timeout_generator( datetime.timedelta(seconds=0.1), itertools.count, 2, + [0.06, 0.1], ), ], ) def test_timeout_generator( + fake_clock: clock.FakeClock, timeout: float, interval: float, interval_multiplier: float, @@ -73,8 +91,9 @@ def test_timeout_generator( types.Callable[..., types.Iterable[types.Any]], ], result: int, + sleeps: types.List[float], ) -> None: - """Stop the sync generator near the configured timeout.""" + """Stop the sync generator at the timeout and scale the interval.""" i = None for i in python_utils.timeout_generator( timeout=timeout, @@ -86,59 +105,97 @@ def test_timeout_generator( assert i is not None assert i == result + assert fake_clock.sleeps == pytest.approx(sleeps) -@pytest.mark.asyncio -async def test_aio_generator_timeout_detector() -> None: - """Raise or exit on per-item and total timeouts.""" - # Make pyright happy - i = None +def test_timeout_generator_real_clock() -> None: + """Keep yielding on the real clock until the timeout has passed.""" + timeout: float = 0.05 + interval: float = 0.01 + start: float = time.perf_counter() + items: types.List[int] = list( + python_utils.timeout_generator(timeout, interval, itertools.count()) + ) + elapsed: float = time.perf_counter() - start + + # A sleep can take longer than requested but never shorter, so these + # hold on any machine. The exact number of items does not. + assert items == list(range(len(items))) + assert len(items) <= timeout / interval + 2 + assert elapsed >= timeout + + +async def stalling_generator() -> types.AsyncGenerator[int, None]: + """Yield 0-4 without waiting, then stall before the next item.""" + for i in range(10): + if i == 5: + await asyncio.sleep(STALL) + yield i + + +def ticking_generator( + fake_clock: clock.FakeClock, +) -> types.AsyncGenerator[int, None]: + """Yield 0-9 and let 0.1 seconds pass on the fake clock for each item.""" async def generator() -> types.AsyncGenerator[int, None]: - """Yield 0-9 with increasing sleeps between items.""" + """Advance the fake clock before every item.""" for i in range(10): - await asyncio.sleep(i / 20.0) + fake_clock.sleep(0.1) yield i + return generator() + + +@pytest.mark.asyncio +async def test_aio_generator_timeout_detector( + fake_clock: clock.FakeClock, +) -> None: + """Raise or exit on per-item and total timeouts.""" + # Make pyright happy + i = None + detector = python_utils.aio_generator_timeout_detector # Test regular timeout with reraise with pytest.raises(asyncio.TimeoutError): - async for i in detector(generator(), 0.25): + async for i in detector(stalling_generator(), 0.05): pass # Test regular timeout with clean exit - async for i in detector(generator(), 0.25, on_timeout=None): + async for i in detector(stalling_generator(), 0.05, on_timeout=None): pass assert i == 4 # Test total timeout with reraise with pytest.raises(asyncio.TimeoutError): - async for i in detector(generator(), total_timeout=0.5): + async for i in detector( + ticking_generator(fake_clock), total_timeout=0.45 + ): pass # Test total timeout with clean exit - async for i in detector(generator(), total_timeout=0.5, on_timeout=None): + async for i in detector( + ticking_generator(fake_clock), total_timeout=0.45, on_timeout=None + ): pass assert i == 4 # Test stop iteration - async for i in detector(generator(), on_timeout=None): + async for i in detector(ticking_generator(fake_clock), on_timeout=None): pass + assert i == 9 + @pytest.mark.asyncio async def test_aio_generator_timeout_detector_decorator_reraise() -> None: """Reraise ``TimeoutError`` on a per-item timeout.""" - # Test regular timeout with reraise - @python_utils.aio_generator_timeout_detector_decorator(timeout=0.05) - async def generator_timeout() -> types.AsyncGenerator[int, None]: - """Yield with increasing delays to trip the timeout.""" - for i in range(10): - await asyncio.sleep(i / 100.0) - yield i + generator_timeout = python_utils.aio_generator_timeout_detector_decorator( + timeout=0.05 + )(stalling_generator) with pytest.raises(asyncio.TimeoutError): async for _ in generator_timeout(): @@ -152,14 +209,9 @@ async def test_aio_generator_timeout_detector_decorator_clean_exit() -> None: i = None # Test regular timeout with clean exit - @python_utils.aio_generator_timeout_detector_decorator( + generator_clean = python_utils.aio_generator_timeout_detector_decorator( timeout=0.05, on_timeout=None - ) - async def generator_clean() -> types.AsyncGenerator[int, None]: - """Yield with increasing delays to trip the timeout.""" - for i in range(10): - await asyncio.sleep(i / 100.0) - yield i + )(stalling_generator) async for i in generator_clean(): pass @@ -168,17 +220,16 @@ async def generator_clean() -> types.AsyncGenerator[int, None]: @pytest.mark.asyncio -async def test_aio_generator_timeout_detector_decorator_reraise_total() -> ( - None -): +async def test_aio_generator_timeout_detector_decorator_reraise_total( + fake_clock: clock.FakeClock, +) -> None: """Reraise ``TimeoutError`` on a total timeout.""" # Test total timeout with reraise - @python_utils.aio_generator_timeout_detector_decorator(total_timeout=0.1) + @python_utils.aio_generator_timeout_detector_decorator(total_timeout=0.45) async def generator_reraise() -> types.AsyncGenerator[int, None]: - """Yield with increasing delays to trip the timeout.""" - for i in range(10): - await asyncio.sleep(i / 100.0) + """Let the fake clock pass the total timeout while yielding.""" + async for i in ticking_generator(fake_clock): yield i with pytest.raises(asyncio.TimeoutError): @@ -187,22 +238,350 @@ async def generator_reraise() -> types.AsyncGenerator[int, None]: @pytest.mark.asyncio -async def test_aio_generator_timeout_detector_decorator_clean_total() -> None: +async def test_aio_generator_timeout_detector_decorator_clean_total( + fake_clock: clock.FakeClock, +) -> None: """Exit cleanly on total timeout when ``on_timeout`` is ``None``.""" # Make pyright happy i = None # Test total timeout with clean exit @python_utils.aio_generator_timeout_detector_decorator( - total_timeout=0.1, on_timeout=None + total_timeout=0.45, on_timeout=None ) async def generator_clean_total() -> types.AsyncGenerator[int, None]: - """Yield with increasing delays to trip the timeout.""" - for i in range(10): - await asyncio.sleep(i / 100.0) + """Let the fake clock pass the total timeout while yielding.""" + async for i in ticking_generator(fake_clock): yield i async for i in generator_clean_total(): pass assert i == 4 + + +@pytest.mark.asyncio +async def test_aio_generator_timeout_detector_sync_callback( + fake_clock: clock.FakeClock, +) -> None: + """Accept an ``on_timeout`` callback that is a plain function.""" + exceptions: types.List[BaseException] = [] + + def on_timeout( + generator: types.AsyncGenerator[int, None], + timeout: types.Optional[types.delta_type], + total_timeout: types.Optional[types.delta_type], + exception: BaseException, + ) -> None: + """Record the timeout and return nothing to await.""" + exceptions.append(exception) + + items: types.List[int] = [ + i + async for i in python_utils.aio_generator_timeout_detector( + ticking_generator(fake_clock), + total_timeout=0.45, + on_timeout=on_timeout, + ) + ] + + assert items == [0, 1, 2, 3, 4] + assert len(exceptions) == 1 + assert isinstance(exceptions[0], asyncio.TimeoutError) + + +@pytest.mark.asyncio +async def test_aio_generator_timeout_detector_async_callback( + fake_clock: clock.FakeClock, +) -> None: + """Await an ``on_timeout`` callback that is a coroutine function.""" + calls: types.List[types.Mapping[str, types.Any]] = [] + + async def on_timeout( + generator: types.AsyncGenerator[int, None], + timeout: types.Optional[types.delta_type], + total_timeout: types.Optional[types.delta_type], + exception: BaseException, + **kwargs: types.Mapping[str, types.Any], + ) -> None: + """Record the extra keyword arguments after a real await.""" + await asyncio.sleep(0) + calls.append(kwargs) + + items: types.List[int] = [ + i + async for i in python_utils.aio_generator_timeout_detector( + ticking_generator(fake_clock), + total_timeout=0.45, + on_timeout=on_timeout, + context={'attempt': 1}, + ) + ] + + assert items == [0, 1, 2, 3, 4] + assert calls == [{'context': {'attempt': 1}}] + + +async def collect_before_stall( + generator: types.AsyncGenerator[int, None], +) -> types.List[int]: + """Collect the items of a detector that has to stop at the stall.""" + items: types.List[int] = [] + + async def collect() -> None: + """Gather every item the detector lets through.""" + item: int + async for item in generator: + items.append(item) + + try: + # Without the guard a detector that waits out the stall would keep + # the test busy for `STALL` seconds before it fails. + await asyncio.wait_for(collect(), STALL / 10) + finally: + await generator.aclose() + + return items + + +@pytest.mark.asyncio +async def test_aio_generator_timeout_detector_item_before_total() -> None: + """Let the item timeout win when it ends before the total timeout.""" + detector: types.AsyncGenerator[int, None] = ( + python_utils.aio_generator_timeout_detector( + stalling_generator(), timeout=STALL / 500, total_timeout=STALL + ) + ) + + with pytest.raises(asyncio.TimeoutError) as exc_info: + await collect_before_stall(detector) + + assert 'Total timeout reached' not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_aio_generator_timeout_detector_total_timeout_finishes() -> None: + """Finish as normal when the generator ends within the total timeout.""" + + async def generator() -> types.AsyncGenerator[int, None]: + """Yield three items without waiting.""" + i: int + for i in range(3): + yield i + + items: types.List[int] = [ + i + async for i in python_utils.aio_generator_timeout_detector( + generator(), total_timeout=STALL + ) + ] + + assert items == [0, 1, 2] + + +@pytest.mark.asyncio +async def test_aio_generator_timeout_detector_generator_timeout() -> None: + """Hand a timeout raised by the generator itself over unchanged.""" + exceptions: types.List[BaseException] = [] + + async def generator() -> types.AsyncGenerator[int, None]: + """Yield one item and then fail with a timeout of its own.""" + yield 0 + raise asyncio.TimeoutError('raised by the generator') + + def on_timeout( + generator: types.AsyncGenerator[int, None], + timeout: types.Optional[types.delta_type], + total_timeout: types.Optional[types.delta_type], + exception: BaseException, + ) -> None: + """Record the exception the detector reports.""" + exceptions.append(exception) + + items: types.List[int] = [ + i + async for i in python_utils.aio_generator_timeout_detector( + generator(), total_timeout=STALL, on_timeout=on_timeout + ) + ] + + assert items == [0] + assert [str(exception) for exception in exceptions] == [ + 'raised by the generator' + ] + + +@pytest.mark.parametrize( + 'timestamp,precision,expected', + [ + (1, datetime.timedelta(milliseconds=100), '0:00:01'), + (60, datetime.timedelta(milliseconds=100), '0:01:00'), + (0.3, datetime.timedelta(milliseconds=100), '0:00:00.300000'), + (1.234, datetime.timedelta(milliseconds=1), '0:00:01.234000'), + (1.239, datetime.timedelta(milliseconds=10), '0:00:01.230000'), + ( + datetime.timedelta(seconds=2), + datetime.timedelta(milliseconds=100), + '0:00:02', + ), + ( + datetime.datetime(2000, 1, 2, 3, 4, 5), + datetime.timedelta(milliseconds=100), + '2000-01-02 03:04:05', + ), + ( + datetime.datetime(2000, 1, 2, 3, 4, 5, 678901), + datetime.timedelta(milliseconds=10), + '2000-01-02 03:04:05.670000', + ), + ], +) +def test_format_time_sub_second_precision( + timestamp: types.timestamp_type, + precision: datetime.timedelta, + expected: str, +) -> None: + """Keep a value that is on the precision grid where it is.""" + assert python_utils.format_time(timestamp, precision) == expected + + +@pytest.mark.parametrize( + 'timestamp', + [float('nan'), 'nan', float('inf'), float('-inf'), 1e20, 10**30], +) +def test_format_time_placeholder_for_impossible_numbers( + timestamp: types.timestamp_type, +) -> None: + """Print the placeholder for a number that is not a duration.""" + assert python_utils.format_time(timestamp) == '--:--:--' + + +def test_format_time_extreme_timedeltas() -> None: + """Format the largest and the smallest timedelta without overflow.""" + largest: str = python_utils.format_time(datetime.timedelta.max) + smallest: str = python_utils.format_time(datetime.timedelta.min) + + assert largest == '999999999 days, 23:59:59' + assert smallest == '-999999999 days, 0:00:00' + + +@pytest.mark.parametrize( + 'delta', + [ + datetime.timedelta(microseconds=1), + datetime.timedelta(microseconds=-1), + datetime.timedelta(seconds=437, microseconds=579262), + datetime.timedelta(days=-5, microseconds=1), + datetime.timedelta(days=1000000, microseconds=999999), + datetime.timedelta.max, + ], +) +def test_timedelta_to_seconds_fraction_precision( + delta: datetime.timedelta, +) -> None: + """Keep the precision of ``total_seconds`` for a fraction of a second.""" + seconds: types.Number = python_utils.timedelta_to_seconds(delta) + + assert isinstance(seconds, float) + assert seconds == delta.total_seconds() + + +@pytest.mark.parametrize( + 'delta,expected', + [ + (datetime.timedelta(0), 0), + (datetime.timedelta(seconds=1), 1), + (datetime.timedelta(seconds=-1), -1), + (datetime.timedelta(days=1), 86400), + (datetime.timedelta(days=999999999, seconds=86399), 86399999999999), + (datetime.timedelta.min, -86399999913600), + ], +) +def test_timedelta_to_seconds_whole_seconds_stay_int( + delta: datetime.timedelta, expected: int +) -> None: + """Return whole seconds as the exact ``int`` they have always been.""" + seconds: types.Number = python_utils.timedelta_to_seconds(delta) + + assert isinstance(seconds, int) + assert seconds == expected + + +def test_timeout_generator_maximum_interval_first_sleep( + fake_clock: clock.FakeClock, +) -> None: + """Hold the first sleep to ``maximum_interval`` as well.""" + items: types.List[str] = list( + python_utils.timeout_generator( + timeout=100, interval=10, iterable='abc', maximum_interval=1 + ) + ) + + assert items == ['a', 'b', 'c'] + assert fake_clock.sleeps == [1, 1, 1] + + +def test_timeout_generator_maximum_interval_keeps_progression( + fake_clock: clock.FakeClock, +) -> None: + """Shorten the first sleep only, the later ones stay as they were.""" + items: types.List[str] = list( + python_utils.timeout_generator( + timeout=100, + interval=3, + iterable='abcd', + interval_multiplier=0.5, + maximum_interval=2, + ) + ) + + assert items == ['a', 'b', 'c', 'd'] + assert fake_clock.sleeps == [2, 1.5, 0.75, 0.375] + + +def test_timeout_generator_maximum_interval_zero( + fake_clock: clock.FakeClock, +) -> None: + """Read a ``maximum_interval`` of zero as no maximum at all.""" + items: types.List[str] = list( + python_utils.timeout_generator( + timeout=100, + interval=1, + iterable='abc', + interval_multiplier=2, + maximum_interval=0, + ) + ) + + assert items == ['a', 'b', 'c'] + assert fake_clock.sleeps == [1, 2, 4] + + +@pytest.mark.asyncio +async def test_aio_timeout_generator_maximum_interval_first_sleep( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Hold the first async sleep to ``maximum_interval`` as well.""" + sleeps: types.List[float] = [] + + async def mock_sleep(delay: float) -> None: + """Record each requested delay instead of sleeping.""" + sleeps.append(delay) + + monkeypatch.setattr(asyncio, 'sleep', mock_sleep) + + async def letters() -> types.AsyncGenerator[str, None]: + """Yield three letters without waiting.""" + letter: str + for letter in 'abc': + yield letter + + items: types.List[str] = [ + item + async for item in python_utils.aio_timeout_generator( + timeout=100, interval=10, iterable=letters, maximum_interval=1 + ) + ] + + assert items == ['a', 'b', 'c'] + assert sleeps == [1, 1, 1] diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..4a22f5b --- /dev/null +++ b/conftest.py @@ -0,0 +1,8 @@ +""" +Load the shared fixtures for the tests and the doctests. + +The fixtures live in the tests package. They are loaded from the repository +root because the doctests in ``python_utils`` need them as well. +""" + +pytest_plugins: tuple[str, ...] = ('_python_utils_tests.clock',) diff --git a/pyproject.toml b/pyproject.toml index c875f23..ce32e4d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ module-root = '' module-name = 'python_utils' # Keep the tests and tox config in the sdist (parity with the old MANIFEST.in) # so downstream packagers can build and test from the source distribution. -source-include = ['_python_utils_tests/**/*.py', 'tox.ini'] +source-include = ['_python_utils_tests/**/*.py', 'conftest.py', 'tox.ini'] [project] name = 'python-utils' diff --git a/python_utils/aio.py b/python_utils/aio.py index c315b9b..2b0854d 100644 --- a/python_utils/aio.py +++ b/python_utils/aio.py @@ -35,7 +35,9 @@ async def acount( start: First value to yield. step: Amount added between successive values. delay: Seconds to ``asyncio.sleep`` between yields. - stop: Exclusive upper bound; ``None`` counts forever. + stop: Exclusive bound, ``None`` counts forever. It is an upper bound + for a positive ``step`` and a lower bound for a negative one, the + way ``range`` reads it. Yields: The successive counter values. @@ -46,7 +48,8 @@ async def acount( [0, 1, 2] """ for item in itertools.count(start, step): # pragma: no branch - if stop is not None and item >= stop: + # A negative step counts down, so `stop` is the lower bound there. + if stop is not None and (item <= stop if step < 0 else item >= stop): break yield item diff --git a/python_utils/decorators.py b/python_utils/decorators.py index 9bd81bf..be22d9c 100644 --- a/python_utils/decorators.py +++ b/python_utils/decorators.py @@ -24,6 +24,11 @@ _T = typing.TypeVar('_T') _P = typing.ParamSpec('_P') +#: Logger of this module. The module-level `logging.debug()` writes to the +#: root logger and installs a handler on it when it has none, which is the +#: application's decision to make. +_logger: logging.Logger = logging.getLogger(__name__) + def set_attributes( **kwargs: typing.Any, @@ -74,6 +79,11 @@ def listify( """ Convert any generator to a list or other type of collection. + A function that returns `None` counts as empty. With `allow_empty` that + gives an empty collection, and with `allow_empty=False` it raises a + `TypeError`. The `None` is all that `allow_empty` looks at: a generator + that yields nothing gives an empty collection either way. + >>> @listify() ... def generator(): ... yield 1 @@ -99,6 +109,13 @@ def listify( ... TypeError: ... `allow_empty` is `False` + >>> @listify(allow_empty=False) + ... def generator_without_items(): + ... yield from () + + >>> generator_without_items() + [] + >>> @listify(collection=set) ... def set_generator(): ... yield 1 @@ -124,6 +141,7 @@ def _listify( ) -> collections.abc.Callable[..., collections.abc.Collection[_T]]: """Materialize ``function``'s result into ``collection``.""" + @functools.wraps(function) def __listify( *args: typing.Any, **kwargs: typing.Any ) -> collections.abc.Collection[_T]: @@ -178,7 +196,7 @@ def __sample(*args: _P.args, **kwargs: _P.kwargs) -> _T | None: if random.random() < sample_rate: return function(*args, **kwargs) else: - logging.debug( + _logger.debug( 'Skipped execution of %r(%r, %r) due to sampling', function, args, @@ -224,13 +242,21 @@ def _wraps_classmethod( wrapper = functools.update_wrapper( wrapper, wrapped, + # The annotations are handled below. Python 3.14 lists + # `__annotate__` here where older versions list + # `__annotations__`, so both names are left out. assigned=tuple( a for a in functools.WRAPPER_ASSIGNMENTS - if a != '__annotations__' + if a not in ('__annotations__', '__annotate__') ), ) - if annotations := getattr(wrapped, '__annotations__', {}): + # A copy, because the dictionary belongs to `wrapped` and that one + # keeps its `self`. + annotations: dict[str, typing.Any] = dict( + getattr(wrapped, '__annotations__', {}) + ) + if annotations: # Drop `self`: the wrapper is a classmethod, so it takes no `self`. annotations.pop('self', None) wrapper.__annotations__ = annotations diff --git a/python_utils/generators.py b/python_utils/generators.py index 20cd9c0..e6571bb 100644 --- a/python_utils/generators.py +++ b/python_utils/generators.py @@ -23,6 +23,41 @@ _T = typing.TypeVar('_T') +async def _cancel(pending: set[asyncio.Future[_T]]) -> None: + """Cancel the items that are still on their way and wait for them. + + Args: + pending: The unfinished futures to cancel. An empty set is fine. + """ + if not pending: + # Nothing is on its way. This also covers a batcher that is closed + # by the garbage collector after its event loop is gone. + return + + for future in pending: + future.cancel() + + # Waiting lets the cancellation reach the generator. A cancelled item is + # of no use to anyone. An error that the generator raises while it is + # being cancelled is reported to the event loop, because raising it here + # would replace the cancellation of the consumer. + outcomes: list[_T | BaseException] = await asyncio.gather( + *pending, return_exceptions=True + ) + for outcome in outcomes: + # A generator that returns when it is cancelled ends its item with + # `StopAsyncIteration`, which is no failure. + if isinstance(outcome, Exception) and not isinstance( + outcome, StopAsyncIteration + ): + asyncio.get_running_loop().call_exception_handler( + { + 'message': 'abatcher source failed while being cancelled', + 'exception': outcome, + } + ) + + async def abatcher( generator: collections.abc.AsyncGenerator[_T, None] | collections.abc.AsyncIterator[_T], @@ -33,6 +68,10 @@ async def abatcher( Asyncio generator wrapper that returns items with a given batch size or interval (whichever is reached first). + Stopping early while the next item is still on its way cancels the + request for it. A source that is an async generator ends at that point. + A source with nothing on its way stays usable. + Args: generator: The async generator or iterator to batch. batch_size (typing.Optional[int], optional): The number of items per @@ -58,45 +97,54 @@ async def abatcher( next_yield: float = time.perf_counter() + interval_s - done: set[asyncio.Task[_T]] - pending: set[asyncio.Task[_T]] = set() - - while True: - try: - done, pending = await asyncio.wait( - pending - or [ - asyncio.create_task( - typing.cast( - collections.abc.Coroutine[None, None, _T], - generator.__anext__(), - ) - ), - ], - timeout=interval_s, - return_when=asyncio.FIRST_COMPLETED, - ) + done: set[asyncio.Future[_T]] + pending: set[asyncio.Future[_T]] = set() - if done: - batch.extend(result.result() for result in done) + try: + while True: + try: + if not pending: + # Store the task before waiting for it. A cancellation + # during the wait still finds it in the `finally` below. + # `ensure_future` takes every awaitable, where + # `create_task` insists on a coroutine. + pending = {asyncio.ensure_future(generator.__anext__())} - except StopAsyncIteration: - if batch: - yield batch + done, pending = await asyncio.wait( + pending, + timeout=interval_s, + return_when=asyncio.FIRST_COMPLETED, + ) - break + if done: + batch.extend(result.result() for result in done) - if batch_size is not None and len(batch) == batch_size: - yield batch - batch = [] + except StopAsyncIteration: + if batch: + yield batch - if interval and batch and time.perf_counter() > next_yield: - yield batch - batch = [] - # Always set the next yield time to the current time. If the - # loop is running slow due to blocking functions we do not - # want to burst too much - next_yield = time.perf_counter() + interval_s + break + + if batch_size is not None and len(batch) == batch_size: + yield batch + batch = [] + # A full batch starts a new interval as well. Otherwise the + # next item is flushed on its own once the old one ran out. + next_yield = time.perf_counter() + interval_s + + if interval and batch and time.perf_counter() > next_yield: + yield batch + batch = [] + # Always set the next yield time to the current time. If the + # loop is running slow due to blocking functions we do not + # want to burst too much + next_yield = time.perf_counter() + interval_s + finally: + # The consumer can stop while the next item is still on its way: an + # early `break`, `aclose()` or a cancellation. Without this the task + # keeps running, takes the next item from the generator and nobody + # ever receives it. + await _cancel(pending) def batcher( diff --git a/python_utils/time.py b/python_utils/time.py index 2ec037f..f4fb5af 100644 --- a/python_utils/time.py +++ b/python_utils/time.py @@ -54,14 +54,15 @@ def timedelta_to_seconds(delta: datetime.timedelta) -> _aliases.Number: >>> '%.6f' % timedelta_to_seconds(timedelta(microseconds=1)) '0.000001' """ + seconds: int = delta.seconds + delta.days * 60 * 60 * 24 + # Only convert to float if needed if delta.microseconds: - total = delta.microseconds * 1e-6 + # Divide the whole microseconds once, the way `total_seconds()` does. + # Adding a float fraction to a large number of seconds loses digits. + return (seconds * 10**6 + delta.microseconds) / 10**6 else: - total = 0 - total += delta.seconds - total += delta.days * 60 * 60 * 24 - return total + return seconds def delta_to_seconds(interval: _aliases.delta_type) -> _aliases.Number: @@ -135,23 +136,22 @@ def format_time( TypeError: Unknown type ... """ - precision_seconds = precision.total_seconds() - if isinstance(timestamp, str): timestamp = float(timestamp) if isinstance(timestamp, (int, float)): try: timestamp = datetime.timedelta(seconds=timestamp) - except OverflowError: # pragma: no cover + except (OverflowError, ValueError): + # Too large for a timedelta, or not a number at all: nan raises + # a ValueError where infinity raises an OverflowError. timestamp = None if isinstance(timestamp, datetime.timedelta): - seconds = timestamp.total_seconds() - # Truncate the number to the given precision - seconds = seconds - (seconds % precision_seconds) - - return str(datetime.timedelta(seconds=seconds)) + # Truncate the number to the given precision. A timedelta counts in + # whole microseconds, which keeps the modulo exact. In float seconds + # `1.0 % 0.1` is just under 0.1 and the result drops a whole step. + return str(timestamp - timestamp % precision) elif isinstance(timestamp, datetime.datetime): # pragma: no cover # Python 2 doesn't have the timestamp method if hasattr(timestamp, 'timestamp'): @@ -159,8 +159,10 @@ def format_time( else: seconds = timedelta_to_seconds(timestamp - epoch) - # Truncate the number to the given precision - seconds = seconds - (seconds % precision_seconds) + # Truncate the number to the given precision, in whole microseconds + # for the same reason as above + since_epoch: datetime.timedelta = datetime.timedelta(seconds=seconds) + seconds = (since_epoch - since_epoch % precision).total_seconds() try: # pragma: no cover dt = datetime.datetime.fromtimestamp(seconds) @@ -223,6 +225,10 @@ def timeout_generator( useful for testing slow APIs so you get a small sample of the data in a reasonable amount of time. + After every sleep the interval is multiplied by `interval_multiplier`. No + sleep is longer than `maximum_interval`, and that includes the first one. + A `maximum_interval` of `None` or `0` means that there is no maximum. + >>> for i in timeout_generator(0.1, 0.06): ... # Put your slow code here ... print(i) @@ -262,7 +268,12 @@ def timeout_generator( if time.perf_counter() >= end: break - time.sleep(float_interval) + # The maximum holds for the first sleep as well. Zero is not a + # maximum here, it means the same as `None`. + if float_maximum_interval: + time.sleep(min(float_interval, float_maximum_interval)) + else: + time.sleep(float_interval) float_interval *= interval_multiplier if float_maximum_interval: @@ -283,10 +294,11 @@ async def aio_timeout_generator( default) until the float_timeout is reached with a configurable float_interval between items. - The interval_exponent automatically increases the float_timeout with each - run. Note that if the float_interval is less than 1, 1/interval_exponent - will be used so the float_interval is always growing. To double the - float_interval with each run, specify 2. + After every sleep the interval is multiplied by `interval_multiplier`. To + double the interval with each run, specify 2. A value below 1 makes the + interval shorter with each run. No sleep is longer than + `maximum_interval`, and that includes the first one. A `maximum_interval` + of `None` or `0` means that there is no maximum. Doctests and asyncio are not friends, so no examples. But this function is effectively the same as the `timeout_generator` but it uses `async for` @@ -317,7 +329,12 @@ async def aio_timeout_generator( if time.perf_counter() >= end: break - await asyncio.sleep(float_interval) + # The maximum holds for the first sleep as well. Zero is not a + # maximum here, it means the same as `None`. + if float_maximum_interval: + await asyncio.sleep(min(float_interval, float_maximum_interval)) + else: + await asyncio.sleep(float_interval) float_interval *= interval_multiplier if float_maximum_interval: # pragma: no branch @@ -344,9 +361,15 @@ async def aio_generator_timeout_detector( This function is used to detect if an asyncio generator has not yielded an element for a set amount of time. + The `timeout` is the time a single element may take. A `timeout` of `None` + or `0` means that there is no timeout per element. The `total_timeout` is + the time all elements together may take. It is checked between elements, + so it does not end the wait for an element that does not arrive. Use + `timeout` for that. + The `on_timeout` argument is called with the `generator`, `timeout`, `total_timeout`, `exception` and the extra `**kwargs` to this function as - arguments. + arguments. It can be a plain function or a coroutine function. If `on_timeout` is not specified, the exception is reraised. If `on_timeout` is `None`, the exception is silently ignored and the generator will finish as normal. @@ -377,13 +400,17 @@ async def aio_generator_timeout_detector( except asyncio.TimeoutError as exception: # noqa: PERF203 if on_timeout is not None: - await on_timeout( + result: typing.Any = on_timeout( generator, timeout, total_timeout, exception, **on_timeout_kwargs, ) + # A coroutine function hands back something to await, a + # plain function has already done its work by now. + if isinstance(result, collections.abc.Awaitable): + await result break except StopAsyncIteration: @@ -411,9 +438,11 @@ def aio_generator_timeout_detector_decorator( """Wrap a generator function with ``aio_generator_timeout_detector``. Args: - timeout: Per-item timeout; if a single yield takes longer, - ``on_timeout`` fires. ``None`` disables the per-item check. - total_timeout: Overall timeout across the whole generator. + timeout: Per-item timeout. If a single yield takes longer, + ``on_timeout`` fires. ``None`` or ``0`` disables the per-item + check. + total_timeout: Overall timeout across the whole generator. It is + checked between items. on_timeout: Callback invoked on a timeout; defaults to re-raising. **on_timeout_kwargs: Extra keyword arguments passed to ``on_timeout``.