/aiostream/stream/time.py

https://github.com/vxgmichel/aiostream · Python · 65 lines · 42 code · 11 blank · 12 comment · 16 complexity · 4475a69f9f365bd7a54a7ec55f16e16f MD5 · raw file

  1. """Time-specific operators."""
  2. import asyncio
  3. from ..aiter_utils import anext
  4. from ..core import operator, streamcontext
  5. __all__ = ['spaceout', 'delay', 'timeout']
  6. async def wait_for(aw, timeout):
  7. task = asyncio.ensure_future(aw)
  8. try:
  9. return await asyncio.wait_for(task, timeout)
  10. finally:
  11. # Python 3.6 compatibility
  12. if not task.done(): # pragma: no cover
  13. task.cancel()
  14. try:
  15. await task
  16. except asyncio.CancelledError:
  17. pass
  18. @operator(pipable=True)
  19. async def spaceout(source, interval):
  20. """Make sure the elements of an asynchronous sequence are separated
  21. in time by the given interval.
  22. """
  23. timeout = 0
  24. loop = asyncio.get_event_loop()
  25. async with streamcontext(source) as streamer:
  26. async for item in streamer:
  27. delta = timeout - loop.time()
  28. delay = delta if delta > 0 else 0
  29. await asyncio.sleep(delay)
  30. yield item
  31. timeout = loop.time() + interval
  32. @operator(pipable=True)
  33. async def timeout(source, timeout):
  34. """Raise a time-out if an element of the asynchronous sequence
  35. takes too long to arrive.
  36. Note: the timeout is not global but specific to each step of
  37. the iteration.
  38. """
  39. async with streamcontext(source) as streamer:
  40. while True:
  41. try:
  42. item = await wait_for(anext(streamer), timeout)
  43. except StopAsyncIteration:
  44. break
  45. else:
  46. yield item
  47. @operator(pipable=True)
  48. async def delay(source, delay):
  49. """Delay the iteration of an asynchronous sequence."""
  50. await asyncio.sleep(delay)
  51. async with streamcontext(source) as streamer:
  52. async for item in streamer:
  53. yield item