From 5c500ee2577fe9cb84b3c6e54216985afc4b485f Mon Sep 17 00:00:00 2001
From: Thomas Grainger <tagrain@gmail.com>
Date: Sun, 15 Oct 2023 09:40:25 -0700
Subject: [PATCH] Restore uvloop.new_event_loop and other missing uvloop
 members to typing (#573)

---
 uvloop/__init__.py  | 118 +++++++++++++++++++++++++-------------------
 uvloop/__init__.pyi |  17 -------
 2 files changed, 68 insertions(+), 67 deletions(-)
 delete mode 100644 uvloop/__init__.pyi

diff --git a/uvloop/__init__.py b/uvloop/__init__.py
index 1fa158a..9bb6592 100644
--- a/uvloop/__init__.py
+++ b/uvloop/__init__.py
@@ -13,6 +13,9 @@ from ._version import __version__  # NOQA
 __all__ = ('new_event_loop', 'install', 'EventLoopPolicy')
 
 
+_T = _typing.TypeVar("_T")
+
+
 class Loop(__BaseLoop, __asyncio.AbstractEventLoop):  # type: ignore[misc]
     pass
 
@@ -34,69 +37,84 @@ def install() -> None:
     __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."""
+if _typing.TYPE_CHECKING:
+    def run(
+        main: _typing.Coroutine[_typing.Any, _typing.Any, _T],
+        *,
+        loop_factory: _typing.Optional[
+            _typing.Callable[[], Loop]
+        ] = new_event_loop,
+        debug: _typing.Optional[bool]=None,
+    ) -> _T:
+        """The preferred way of running a coroutine with uvloop."""
+else:
+    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
+        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]
+        vi = _sys.version_info[:2]
 
-    if vi <= (3, 10):
-        # Copied from python/cpython
+        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 __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))
+            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:
+            loop = loop_factory()
             try:
-                _cancel_all_tasks(loop)
-                loop.run_until_complete(loop.shutdown_asyncgens())
-                if hasattr(loop, 'shutdown_default_executor'):
-                    loop.run_until_complete(loop.shutdown_default_executor())
+                __asyncio.set_event_loop(loop)
+                if debug is not None:
+                    loop.set_debug(debug)
+                return loop.run_until_complete(wrapper())
             finally:
-                __asyncio.set_event_loop(None)
-                loop.close()
+                try:
+                    _cancel_all_tasks(loop)
+                    loop.run_until_complete(loop.shutdown_asyncgens())
+                    if hasattr(loop, 'shutdown_default_executor'):
+                        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")
+        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())
+            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
-        )
+        else:
+            assert vi >= (3, 12)
+            return __asyncio.run(
+                wrapper(),
+                loop_factory=loop_factory,
+                debug=debug,
+                **run_kwargs
+            )
 
 
-def _cancel_all_tasks(loop):
+def _cancel_all_tasks(loop: __asyncio.AbstractEventLoop) -> None:
     # Copied from python/cpython
 
     to_cancel = __asyncio.all_tasks(loop)
diff --git a/uvloop/__init__.pyi b/uvloop/__init__.pyi
deleted file mode 100644
index 06194d4..0000000
--- a/uvloop/__init__.pyi
+++ /dev/null
@@ -1,17 +0,0 @@
-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: ...