__init__.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. import asyncio as __asyncio
  2. import typing as _typing
  3. import sys as _sys
  4. import warnings as _warnings
  5. from asyncio.events import BaseDefaultEventLoopPolicy as __BasePolicy
  6. from . import includes as __includes # NOQA
  7. from .loop import Loop as __BaseLoop # NOQA
  8. from ._version import __version__ # NOQA
  9. __all__ = ('new_event_loop', 'install', 'EventLoopPolicy')
  10. _T = _typing.TypeVar("_T")
  11. class Loop(__BaseLoop, __asyncio.AbstractEventLoop): # type: ignore[misc]
  12. pass
  13. def new_event_loop() -> Loop:
  14. """Return a new event loop."""
  15. return Loop()
  16. def install() -> None:
  17. """A helper function to install uvloop policy."""
  18. if _sys.version_info[:2] >= (3, 12):
  19. _warnings.warn(
  20. 'uvloop.install() is deprecated in favor of uvloop.run() '
  21. 'starting with Python 3.12.',
  22. DeprecationWarning,
  23. stacklevel=1,
  24. )
  25. __asyncio.set_event_loop_policy(EventLoopPolicy())
  26. if _typing.TYPE_CHECKING:
  27. def run(
  28. main: _typing.Coroutine[_typing.Any, _typing.Any, _T],
  29. *,
  30. loop_factory: _typing.Optional[
  31. _typing.Callable[[], Loop]
  32. ] = new_event_loop,
  33. debug: _typing.Optional[bool]=None,
  34. ) -> _T:
  35. """The preferred way of running a coroutine with uvloop."""
  36. else:
  37. def run(main, *, loop_factory=new_event_loop, debug=None, **run_kwargs):
  38. """The preferred way of running a coroutine with uvloop."""
  39. async def wrapper():
  40. # If `loop_factory` is provided we want it to return
  41. # either uvloop.Loop or a subtype of it, assuming the user
  42. # is using `uvloop.run()` intentionally.
  43. loop = __asyncio._get_running_loop()
  44. if not isinstance(loop, Loop):
  45. raise TypeError('uvloop.run() uses a non-uvloop event loop')
  46. return await main
  47. vi = _sys.version_info[:2]
  48. if vi <= (3, 10):
  49. # Copied from python/cpython
  50. if __asyncio._get_running_loop() is not None:
  51. raise RuntimeError(
  52. "asyncio.run() cannot be called from a running event loop")
  53. if not __asyncio.iscoroutine(main):
  54. raise ValueError(
  55. "a coroutine was expected, got {!r}".format(main)
  56. )
  57. loop = loop_factory()
  58. try:
  59. __asyncio.set_event_loop(loop)
  60. if debug is not None:
  61. loop.set_debug(debug)
  62. return loop.run_until_complete(wrapper())
  63. finally:
  64. try:
  65. _cancel_all_tasks(loop)
  66. loop.run_until_complete(loop.shutdown_asyncgens())
  67. if hasattr(loop, 'shutdown_default_executor'):
  68. loop.run_until_complete(
  69. loop.shutdown_default_executor()
  70. )
  71. finally:
  72. __asyncio.set_event_loop(None)
  73. loop.close()
  74. elif vi == (3, 11):
  75. if __asyncio._get_running_loop() is not None:
  76. raise RuntimeError(
  77. "asyncio.run() cannot be called from a running event loop")
  78. with __asyncio.Runner(
  79. loop_factory=loop_factory,
  80. debug=debug,
  81. **run_kwargs
  82. ) as runner:
  83. return runner.run(wrapper())
  84. else:
  85. assert vi >= (3, 12)
  86. return __asyncio.run(
  87. wrapper(),
  88. loop_factory=loop_factory,
  89. debug=debug,
  90. **run_kwargs
  91. )
  92. def _cancel_all_tasks(loop: __asyncio.AbstractEventLoop) -> None:
  93. # Copied from python/cpython
  94. to_cancel = __asyncio.all_tasks(loop)
  95. if not to_cancel:
  96. return
  97. for task in to_cancel:
  98. task.cancel()
  99. loop.run_until_complete(
  100. __asyncio.gather(*to_cancel, return_exceptions=True)
  101. )
  102. for task in to_cancel:
  103. if task.cancelled():
  104. continue
  105. if task.exception() is not None:
  106. loop.call_exception_handler({
  107. 'message': 'unhandled exception during asyncio.run() shutdown',
  108. 'exception': task.exception(),
  109. 'task': task,
  110. })
  111. class EventLoopPolicy(__BasePolicy):
  112. """Event loop policy.
  113. The preferred way to make your application use uvloop:
  114. >>> import asyncio
  115. >>> import uvloop
  116. >>> asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
  117. >>> asyncio.get_event_loop()
  118. <uvloop.Loop running=False closed=False debug=False>
  119. """
  120. def _loop_factory(self) -> Loop:
  121. return new_event_loop()
  122. if _typing.TYPE_CHECKING:
  123. # EventLoopPolicy doesn't implement these, but since they are marked
  124. # as abstract in typeshed, we have to put them in so mypy thinks
  125. # the base methods are overridden. This is the same approach taken
  126. # for the Windows event loop policy classes in typeshed.
  127. def get_child_watcher(self) -> _typing.NoReturn:
  128. ...
  129. def set_child_watcher(
  130. self, watcher: _typing.Any
  131. ) -> _typing.NoReturn:
  132. ...