You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

115 lines
3.6 KiB

4 years ago
  1. import asyncio
  2. import sys
  3. from types import TracebackType
  4. from typing import Optional, Type, Any # noqa
  5. __version__ = '3.0.1'
  6. PY_37 = sys.version_info >= (3, 7)
  7. class timeout:
  8. """timeout context manager.
  9. Useful in cases when you want to apply timeout logic around block
  10. of code or in cases when asyncio.wait_for is not suitable. For example:
  11. >>> with timeout(0.001):
  12. ... async with aiohttp.get('https://github.com') as r:
  13. ... await r.text()
  14. timeout - value in seconds or None to disable timeout logic
  15. loop - asyncio compatible event loop
  16. """
  17. def __init__(self, timeout: Optional[float],
  18. *, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
  19. self._timeout = timeout
  20. if loop is None:
  21. loop = asyncio.get_event_loop()
  22. self._loop = loop
  23. self._task = None # type: Optional[asyncio.Task[Any]]
  24. self._cancelled = False
  25. self._cancel_handler = None # type: Optional[asyncio.Handle]
  26. self._cancel_at = None # type: Optional[float]
  27. def __enter__(self) -> 'timeout':
  28. return self._do_enter()
  29. def __exit__(self,
  30. exc_type: Type[BaseException],
  31. exc_val: BaseException,
  32. exc_tb: TracebackType) -> Optional[bool]:
  33. self._do_exit(exc_type)
  34. return None
  35. async def __aenter__(self) -> 'timeout':
  36. return self._do_enter()
  37. async def __aexit__(self,
  38. exc_type: Type[BaseException],
  39. exc_val: BaseException,
  40. exc_tb: TracebackType) -> None:
  41. self._do_exit(exc_type)
  42. @property
  43. def expired(self) -> bool:
  44. return self._cancelled
  45. @property
  46. def remaining(self) -> Optional[float]:
  47. if self._cancel_at is not None:
  48. return max(self._cancel_at - self._loop.time(), 0.0)
  49. else:
  50. return None
  51. def _do_enter(self) -> 'timeout':
  52. # Support Tornado 5- without timeout
  53. # Details: https://github.com/python/asyncio/issues/392
  54. if self._timeout is None:
  55. return self
  56. self._task = current_task(self._loop)
  57. if self._task is None:
  58. raise RuntimeError('Timeout context manager should be used '
  59. 'inside a task')
  60. if self._timeout <= 0:
  61. self._loop.call_soon(self._cancel_task)
  62. return self
  63. self._cancel_at = self._loop.time() + self._timeout
  64. self._cancel_handler = self._loop.call_at(
  65. self._cancel_at, self._cancel_task)
  66. return self
  67. def _do_exit(self, exc_type: Type[BaseException]) -> None:
  68. if exc_type is asyncio.CancelledError and self._cancelled:
  69. self._cancel_handler = None
  70. self._task = None
  71. raise asyncio.TimeoutError
  72. if self._timeout is not None and self._cancel_handler is not None:
  73. self._cancel_handler.cancel()
  74. self._cancel_handler = None
  75. self._task = None
  76. return None
  77. def _cancel_task(self) -> None:
  78. if self._task is not None:
  79. self._task.cancel()
  80. self._cancelled = True
  81. def current_task(loop: asyncio.AbstractEventLoop) -> 'asyncio.Task[Any]':
  82. if PY_37:
  83. task = asyncio.current_task(loop=loop) # type: ignore
  84. else:
  85. task = asyncio.Task.current_task(loop=loop)
  86. if task is None:
  87. # this should be removed, tokio must use register_task and family API
  88. if hasattr(loop, 'current_task'):
  89. task = loop.current_task() # type: ignore
  90. return task