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.

timeout.py 3.8KB

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