__init__.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233
  1. import asyncio as __asyncio
  2. import typing as _typing
  3. import sys as _sys
  4. import warnings as _warnings
  5. from . import includes as __includes # NOQA
  6. from .loop import Loop as __BaseLoop # NOQA
  7. from ._version import __version__ # NOQA
  8. __all__: _typing.Tuple[str, ...] = ('new_event_loop', 'run')
  9. _AbstractEventLoop = __asyncio.AbstractEventLoop
  10. _T = _typing.TypeVar("_T")
  11. class Loop(__BaseLoop, _AbstractEventLoop): # type: ignore[misc]
  12. pass
  13. def new_event_loop() -> Loop:
  14. """Return a new event loop."""
  15. return Loop()
  16. if _typing.TYPE_CHECKING:
  17. def run(
  18. main: _typing.Coroutine[_typing.Any, _typing.Any, _T],
  19. *,
  20. loop_factory: _typing.Optional[
  21. _typing.Callable[[], Loop]
  22. ] = new_event_loop,
  23. debug: _typing.Optional[bool]=None,
  24. ) -> _T:
  25. """The preferred way of running a coroutine with uvloop."""
  26. else:
  27. def run(main, *, loop_factory=new_event_loop, debug=None, **run_kwargs):
  28. """The preferred way of running a coroutine with uvloop."""
  29. async def wrapper():
  30. # If `loop_factory` is provided we want it to return
  31. # either uvloop.Loop or a subtype of it, assuming the user
  32. # is using `uvloop.run()` intentionally.
  33. loop = __asyncio._get_running_loop()
  34. if not isinstance(loop, Loop):
  35. raise TypeError('uvloop.run() uses a non-uvloop event loop')
  36. return await main
  37. vi = _sys.version_info[:2]
  38. if vi <= (3, 10):
  39. # Copied from python/cpython
  40. if __asyncio._get_running_loop() is not None:
  41. raise RuntimeError(
  42. "asyncio.run() cannot be called from a running event loop")
  43. if not __asyncio.iscoroutine(main):
  44. raise ValueError(
  45. "a coroutine was expected, got {!r}".format(main)
  46. )
  47. loop = loop_factory()
  48. try:
  49. __asyncio.set_event_loop(loop)
  50. if debug is not None:
  51. loop.set_debug(debug)
  52. return loop.run_until_complete(wrapper())
  53. finally:
  54. try:
  55. _cancel_all_tasks(loop)
  56. loop.run_until_complete(loop.shutdown_asyncgens())
  57. if hasattr(loop, 'shutdown_default_executor'):
  58. loop.run_until_complete(
  59. loop.shutdown_default_executor()
  60. )
  61. finally:
  62. __asyncio.set_event_loop(None)
  63. loop.close()
  64. elif vi == (3, 11):
  65. if __asyncio._get_running_loop() is not None:
  66. raise RuntimeError(
  67. "asyncio.run() cannot be called from a running event loop")
  68. with __asyncio.Runner(
  69. loop_factory=loop_factory,
  70. debug=debug,
  71. **run_kwargs
  72. ) as runner:
  73. return runner.run(wrapper())
  74. else:
  75. assert vi >= (3, 12)
  76. return __asyncio.run(
  77. wrapper(),
  78. loop_factory=loop_factory,
  79. debug=debug,
  80. **run_kwargs
  81. )
  82. def _cancel_all_tasks(loop: _AbstractEventLoop) -> None:
  83. # Copied from python/cpython
  84. to_cancel = __asyncio.all_tasks(loop)
  85. if not to_cancel:
  86. return
  87. for task in to_cancel:
  88. task.cancel()
  89. loop.run_until_complete(
  90. __asyncio.gather(*to_cancel, return_exceptions=True)
  91. )
  92. for task in to_cancel:
  93. if task.cancelled():
  94. continue
  95. if task.exception() is not None:
  96. loop.call_exception_handler({
  97. 'message': 'unhandled exception during asyncio.run() shutdown',
  98. 'exception': task.exception(),
  99. 'task': task,
  100. })
  101. _deprecated_names = ('install', 'EventLoopPolicy')
  102. if _sys.version_info[:2] < (3, 16):
  103. __all__ += _deprecated_names
  104. def __getattr__(name: str) -> _typing.Any:
  105. if name not in _deprecated_names:
  106. raise AttributeError(f"module 'uvloop' has no attribute '{name}'")
  107. elif _sys.version_info[:2] >= (3, 16):
  108. raise AttributeError(
  109. f"module 'uvloop' has no attribute '{name}' "
  110. f"(it was removed in Python 3.16, use uvloop.run() instead)"
  111. )
  112. import threading
  113. def install() -> None:
  114. """A helper function to install uvloop policy.
  115. This function is deprecated and will be removed in Python 3.16.
  116. Use `uvloop.run()` instead.
  117. """
  118. if _sys.version_info[:2] >= (3, 12):
  119. _warnings.warn(
  120. 'uvloop.install() is deprecated in favor of uvloop.run() '
  121. 'starting with Python 3.12.',
  122. DeprecationWarning,
  123. stacklevel=1,
  124. )
  125. __asyncio.set_event_loop_policy(EventLoopPolicy())
  126. class EventLoopPolicy(
  127. # This is to avoid a mypy error about AbstractEventLoopPolicy
  128. getattr(__asyncio, 'AbstractEventLoopPolicy') # type: ignore[misc]
  129. ):
  130. """Event loop policy for uvloop.
  131. This class is deprecated and will be removed in Python 3.16.
  132. Use `uvloop.run()` instead.
  133. >>> import asyncio
  134. >>> import uvloop
  135. >>> asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
  136. >>> asyncio.get_event_loop()
  137. <uvloop.Loop running=False closed=False debug=False>
  138. """
  139. def _loop_factory(self) -> Loop:
  140. return new_event_loop()
  141. if _typing.TYPE_CHECKING:
  142. # EventLoopPolicy doesn't implement these, but since they are
  143. # marked as abstract in typeshed, we have to put them in so mypy
  144. # thinks the base methods are overridden. This is the same approach
  145. # taken for the Windows event loop policy classes in typeshed.
  146. def get_child_watcher(self) -> _typing.NoReturn:
  147. ...
  148. def set_child_watcher(
  149. self, watcher: _typing.Any
  150. ) -> _typing.NoReturn:
  151. ...
  152. class _Local(threading.local):
  153. _loop: _typing.Optional[_AbstractEventLoop] = None
  154. def __init__(self) -> None:
  155. self._local = self._Local()
  156. def get_event_loop(self) -> _AbstractEventLoop:
  157. """Get the event loop for the current context.
  158. Returns an instance of EventLoop or raises an exception.
  159. """
  160. if self._local._loop is None:
  161. raise RuntimeError(
  162. 'There is no current event loop in thread %r.'
  163. % threading.current_thread().name
  164. )
  165. return self._local._loop
  166. def set_event_loop(
  167. self, loop: _typing.Optional[_AbstractEventLoop]
  168. ) -> None:
  169. """Set the event loop."""
  170. if loop is not None and not isinstance(loop, _AbstractEventLoop):
  171. raise TypeError(
  172. f"loop must be an instance of AbstractEventLoop or None, "
  173. f"not '{type(loop).__name__}'"
  174. )
  175. self._local._loop = loop
  176. def new_event_loop(self) -> Loop:
  177. """Create a new event loop.
  178. You must call set_event_loop() to make this the current event loop.
  179. """
  180. return self._loop_factory()
  181. globals()['install'] = install
  182. globals()['EventLoopPolicy'] = EventLoopPolicy
  183. return globals()[name]