mirror of
https://github.com/PiBrewing/craftbeerpi4.git
synced 2024-11-23 07:28:13 +01:00
116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
|
import asyncio
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from types import TracebackType
|
||
|
|
from typing import Optional, Type, Any # noqa
|
||
|
|
|
||
|
|
|
||
|
|
__version__ = '3.0.1'
|
||
|
|
|
||
|
|
PY_37 = sys.version_info >= (3, 7)
|
||
|
|
|
||
|
|
|
||
|
|
class timeout:
|
||
|
|
"""timeout context manager.
|
||
|
|
|
||
|
|
Useful in cases when you want to apply timeout logic around block
|
||
|
|
of code or in cases when asyncio.wait_for is not suitable. For example:
|
||
|
|
|
||
|
|
>>> with timeout(0.001):
|
||
|
|
... async with aiohttp.get('https://github.com') as r:
|
||
|
|
... await r.text()
|
||
|
|
|
||
|
|
|
||
|
|
timeout - value in seconds or None to disable timeout logic
|
||
|
|
loop - asyncio compatible event loop
|
||
|
|
"""
|
||
|
|
def __init__(self, timeout: Optional[float],
|
||
|
|
*, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
|
||
|
|
self._timeout = timeout
|
||
|
|
if loop is None:
|
||
|
|
loop = asyncio.get_event_loop()
|
||
|
|
self._loop = loop
|
||
|
|
self._task = None # type: Optional[asyncio.Task[Any]]
|
||
|
|
self._cancelled = False
|
||
|
|
self._cancel_handler = None # type: Optional[asyncio.Handle]
|
||
|
|
self._cancel_at = None # type: Optional[float]
|
||
|
|
|
||
|
|
def __enter__(self) -> 'timeout':
|
||
|
|
return self._do_enter()
|
||
|
|
|
||
|
|
def __exit__(self,
|
||
|
|
exc_type: Type[BaseException],
|
||
|
|
exc_val: BaseException,
|
||
|
|
exc_tb: TracebackType) -> Optional[bool]:
|
||
|
|
self._do_exit(exc_type)
|
||
|
|
return None
|
||
|
|
|
||
|
|
async def __aenter__(self) -> 'timeout':
|
||
|
|
return self._do_enter()
|
||
|
|
|
||
|
|
async def __aexit__(self,
|
||
|
|
exc_type: Type[BaseException],
|
||
|
|
exc_val: BaseException,
|
||
|
|
exc_tb: TracebackType) -> None:
|
||
|
|
self._do_exit(exc_type)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def expired(self) -> bool:
|
||
|
|
return self._cancelled
|
||
|
|
|
||
|
|
@property
|
||
|
|
def remaining(self) -> Optional[float]:
|
||
|
|
if self._cancel_at is not None:
|
||
|
|
return max(self._cancel_at - self._loop.time(), 0.0)
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def _do_enter(self) -> 'timeout':
|
||
|
|
# Support Tornado 5- without timeout
|
||
|
|
# Details: https://github.com/python/asyncio/issues/392
|
||
|
|
if self._timeout is None:
|
||
|
|
return self
|
||
|
|
|
||
|
|
self._task = current_task(self._loop)
|
||
|
|
if self._task is None:
|
||
|
|
raise RuntimeError('Timeout context manager should be used '
|
||
|
|
'inside a task')
|
||
|
|
|
||
|
|
if self._timeout <= 0:
|
||
|
|
self._loop.call_soon(self._cancel_task)
|
||
|
|
return self
|
||
|
|
|
||
|
|
self._cancel_at = self._loop.time() + self._timeout
|
||
|
|
self._cancel_handler = self._loop.call_at(
|
||
|
|
self._cancel_at, self._cancel_task)
|
||
|
|
return self
|
||
|
|
|
||
|
|
def _do_exit(self, exc_type: Type[BaseException]) -> None:
|
||
|
|
if exc_type is asyncio.CancelledError and self._cancelled:
|
||
|
|
self._cancel_handler = None
|
||
|
|
self._task = None
|
||
|
|
raise asyncio.TimeoutError
|
||
|
|
if self._timeout is not None and self._cancel_handler is not None:
|
||
|
|
self._cancel_handler.cancel()
|
||
|
|
self._cancel_handler = None
|
||
|
|
self._task = None
|
||
|
|
return None
|
||
|
|
|
||
|
|
def _cancel_task(self) -> None:
|
||
|
|
if self._task is not None:
|
||
|
|
self._task.cancel()
|
||
|
|
self._cancelled = True
|
||
|
|
|
||
|
|
|
||
|
|
def current_task(loop: asyncio.AbstractEventLoop) -> 'asyncio.Task[Any]':
|
||
|
|
if PY_37:
|
||
|
|
task = asyncio.current_task(loop=loop) # type: ignore
|
||
|
|
else:
|
||
|
|
task = asyncio.Task.current_task(loop=loop)
|
||
|
|
if task is None:
|
||
|
|
# this should be removed, tokio must use register_task and family API
|
||
|
|
if hasattr(loop, 'current_task'):
|
||
|
|
task = loop.current_task() # type: ignore
|
||
|
|
|
||
|
|
return task
|