Implement uvloop.run()

Author: 1st1

Diff for: README.rst

@@ -53,6 +53,27 @@ uvloop with::
 Using uvloop
 ------------
 
+As of uvloop 0.18, the preferred way of using it is via the
+``uvloop.run()`` helper function:
+
+
+.. code:: python
+
+    import uvloop
+
+    async def main():
+        # Main entry-point.
+        ...
+
+    uvloop.run(main())
+
+``uvloop.run()`` works by simply configuring ``asyncio.run()``
+to use uvloop, passing all of the arguments to it, such as ``debug``,
+e.g. ``uvloop.run(main(), debug=True)``.
+
+With Python 3.11 and earlier the following alternative
+snippet can be used:
+
 .. code:: python
 
     import asyncio

Diff for: tests/test_dealloc.py

@@ -30,14 +30,12 @@ def test_dealloc_1(self):
         async def test():
             prog = '''\
 import uvloop
-import asyncio
 
 async def foo():
     return 42
 
 def main():
-    uvloop.install()
-    loop = asyncio.get_event_loop()
+    loop = uvloop.new_event_loop()
     loop.set_debug(True)
     loop.run_until_complete(foo())
     # Do not close the loop on purpose: let __dealloc__ methods run.

Diff for: tests/test_process.py

@@ -851,8 +851,7 @@ async def test():
                 )
                 await proc.communicate()
 
-            uvloop.install()
-            asyncio.run(test())
+            uvloop.run(test())
 
             stdin, stdout, stderr = dups
             (r, w), = pipes

Diff for: tests/test_runner.py

@@ -0,0 +1,39 @@
+import asyncio
+import unittest
+import uvloop
+
+
+class TestSourceCode(unittest.TestCase):
+
+    def test_uvloop_run_1(self):
+        CNT = 0
+
+        async def main():
+            nonlocal CNT
+            CNT += 1
+
+            loop = asyncio.get_running_loop()
+
+            self.assertTrue(isinstance(loop, uvloop.Loop))
+            self.assertTrue(loop.get_debug())
+
+            return 'done'
+
+        result = uvloop.run(main(), debug=True)
+
+        self.assertEqual(result, 'done')
+        self.assertEqual(CNT, 1)
+
+    def test_uvloop_run_2(self):
+
+        async def main():
+            pass
+
+        coro = main()
+        with self.assertRaisesRegex(TypeError, ' a non-uvloop event loop'):
+            uvloop.run(
+                coro,
+                loop_factory=asyncio.DefaultEventLoopPolicy().new_event_loop,
+            )
+
+        coro.close()

Diff for: uvloop/__init__.py

@@ -1,5 +1,7 @@
 import asyncio as __asyncio
 import typing as _typing
+import sys as _sys
+import warnings as _warnings
 
 from asyncio.events import BaseDefaultEventLoopPolicy as __BasePolicy
 
@@ -22,9 +24,102 @@ def new_event_loop() -> Loop:
 
 def install() -> None:
     """A helper function to install uvloop policy."""
+    if _sys.version_info[:2] >= (3, 12):
+        _warnings.warn(
+            'uvloop.install() is deprecated in favor of uvloop.run() '
+            'starting with Python 3.12.',
+            DeprecationWarning,
+            stacklevel=1,
+        )
     __asyncio.set_event_loop_policy(EventLoopPolicy())
 
 
+def run(main, *, loop_factory=new_event_loop, debug=None, **run_kwargs):
+    """The preferred way of running a coroutine with uvloop."""
+
+    async def wrapper():
+        # If `loop_factory` is provided we want it to return
+        # either uvloop.Loop or a subtype of it, assuming the user
+        # is using `uvloop.run()` intentionally.
+        loop = __asyncio._get_running_loop()
+        if not isinstance(loop, Loop):
+            raise TypeError('uvloop.run() uses a non-uvloop event loop')
+        return await main
+
+    vi = _sys.version_info[:2]
+
+    if vi <= (3, 10):
+        # Copied from python/cpython
+
+        if __asyncio._get_running_loop() is not None:
+            raise RuntimeError(
+                "asyncio.run() cannot be called from a running event loop")
+
+        if not __asyncio.iscoroutine(main):
+            raise ValueError("a coroutine was expected, got {!r}".format(main))
+
+        loop = loop_factory()
+        try:
+            __asyncio.set_event_loop(loop)
+            if debug is not None:
+                loop.set_debug(debug)
+            return loop.run_until_complete(wrapper())
+        finally:
+            try:
+                _cancel_all_tasks(loop)
+                loop.run_until_complete(loop.shutdown_asyncgens())
+                loop.run_until_complete(loop.shutdown_default_executor())
+            finally:
+                __asyncio.set_event_loop(None)
+                loop.close()
+
+    elif vi == (3, 11):
+        if __asyncio._get_running_loop() is not None:
+            raise RuntimeError(
+                "asyncio.run() cannot be called from a running event loop")
+
+        with __asyncio.Runner(
+            loop_factory=loop_factory,
+            debug=debug,
+            **run_kwargs
+        ) as runner:
+            return runner.run(wrapper())
+
+    else:
+        assert vi >= (3, 12)
+        return __asyncio.run(
+            wrapper(),
+            loop_factory=loop_factory,
+            debug=debug,
+            **run_kwargs
+        )
+
+
+def _cancel_all_tasks(loop):
+    # Copied from python/cpython
+
+    to_cancel = __asyncio.all_tasks(loop)
+    if not to_cancel:
+        return
+
+    for task in to_cancel:
+        task.cancel()
+
+    loop.run_until_complete(
+        __asyncio.gather(*to_cancel, return_exceptions=True)
+    )
+
+    for task in to_cancel:
+        if task.cancelled():
+            continue
+        if task.exception() is not None:
+            loop.call_exception_handler({
+                'message': 'unhandled exception during asyncio.run() shutdown',
+                'exception': task.exception(),
+                'task': task,
+            })
+
+
 class EventLoopPolicy(__BasePolicy):
     """Event loop policy.
 

Diff for: uvloop/__init__.pyi

@@ -0,0 +1,17 @@
+import sys
+
+from asyncio import AbstractEventLoop
+from collections.abc import Callable, Coroutine
+from contextvars import Context
+from typing import Any, TypeVar
+
+
+_T = TypeVar('_T')
+
+
+def run(
+    main: Coroutine[Any, Any, _T],
+    *,
+    debug: bool | None = ...,
+    loop_factory: Callable[[], AbstractEventLoop] | None = ...,
+) -> _T: ...