wehub-resource-sync
70 lines
1.6 KiB
Python
70 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Generator
|
|
from typing import Any
|
|
|
|
|
|
def _finish_fut(fut: asyncio.Future[Any]) -> None:
|
|
if fut.cancelled():
|
|
return
|
|
|
|
fut.set_result(None)
|
|
|
|
|
|
class SleepFinished(Exception):
|
|
pass
|
|
|
|
|
|
class Sleep:
|
|
"""Same as asyncio.sleep except it is resettable"""
|
|
|
|
def __init__(self, delay: float) -> None:
|
|
self._delay = delay
|
|
self._handler: asyncio.TimerHandle | None = None
|
|
|
|
def reset(self, new_delay: float | None = None) -> None:
|
|
if new_delay is None:
|
|
new_delay = self._delay
|
|
|
|
self._delay = new_delay
|
|
|
|
if self._handler is None:
|
|
return
|
|
|
|
if self._handler.cancelled() or self._fut.done():
|
|
raise SleepFinished
|
|
|
|
self._handler.cancel()
|
|
loop = asyncio.get_event_loop()
|
|
self._handler = loop.call_later(new_delay, _finish_fut, self._fut)
|
|
|
|
def cancel(self) -> None:
|
|
if self._handler is None:
|
|
return
|
|
|
|
self._handler.cancel()
|
|
self._fut.cancel()
|
|
|
|
async def _sleep(self) -> None:
|
|
if self._delay <= 0:
|
|
self._fut = asyncio.Future[None]()
|
|
self._fut.set_result(None)
|
|
return
|
|
|
|
loop = asyncio.get_event_loop()
|
|
self._fut = loop.create_future()
|
|
self._handler = loop.call_later(self._delay, _finish_fut, self._fut)
|
|
|
|
try:
|
|
await asyncio.shield(self._fut)
|
|
finally:
|
|
self._handler.cancel()
|
|
|
|
def __await__(self) -> Generator[Any, Any, None]:
|
|
return self._sleep().__await__()
|
|
|
|
|
|
def sleep(delay: float) -> Sleep:
|
|
return Sleep(delay)
|