From ec9c6ce650a1ec9042a2d393036b59f331d85ca5 Mon Sep 17 00:00:00 2001
From: Thomas Grainger <tagrain@gmail.com>
Date: Sat, 14 Oct 2023 12:14:13 +0100
Subject: [PATCH] restore missing uvloop members to uvloop typing

---
 uvloop/__init__.py  | 24 +++++++++++++++++++++---
 uvloop/__init__.pyi | 17 -----------------
 2 files changed, 21 insertions(+), 20 deletions(-)
 delete mode 100644 uvloop/__init__.pyi

diff --git a/uvloop/__init__.py b/uvloop/__init__.py
index 1fa158a6..50165982 100644
--- a/uvloop/__init__.py
+++ b/uvloop/__init__.py
@@ -13,6 +13,9 @@
 __all__ = ('new_event_loop', 'install', 'EventLoopPolicy')
 
 
+_T = _typing.TypeVar("_T")
+
+
 class Loop(__BaseLoop, __asyncio.AbstractEventLoop):  # type: ignore[misc]
     pass
 
@@ -34,10 +37,25 @@ def install() -> None:
     __asyncio.set_event_loop_policy(EventLoopPolicy())
 
 
-def run(main, *, loop_factory=new_event_loop, debug=None, **run_kwargs):
+@_typing.overload
+def run(
+    main: _typing.Coroutine[_typing.Any, _typing.Any, _T],
+    *,
+    debug: _typing.Optional[bool] = ...,
+    loop_factory: _typing.Optional[_typing.Callable[[], Loop]] = ...,
+) -> _T: ...
+
+
+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,
+    **run_kwargs: object
+) -> _T:
     """The preferred way of running a coroutine with uvloop."""
 
-    async def wrapper():
+    async def wrapper() -> _T:
         # 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.
@@ -96,7 +114,7 @@ async def wrapper():
         )
 
 
-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 06194d47..00000000
--- 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: ...