Skip to content

Commit

Permalink
fix typing inside run
Browse files Browse the repository at this point in the history
  • Loading branch information
graingert committed Oct 14, 2023
1 parent ec9c6ce commit d887fa4
Showing 1 changed file with 34 additions and 47 deletions.
81 changes: 34 additions & 47 deletions uvloop/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import asyncio as __asyncio
import typing as _typing
import sys as _sys
import sys
import warnings as _warnings

from asyncio.events import BaseDefaultEventLoopPolicy as __BasePolicy
Expand All @@ -27,7 +27,7 @@ def new_event_loop() -> Loop:

def install() -> None:
"""A helper function to install uvloop policy."""
if _sys.version_info[:2] >= (3, 12):
if sys.version_info[:2] >= (3, 12):
_warnings.warn(
'uvloop.install() is deprecated in favor of uvloop.run() '
'starting with Python 3.12.',
Expand All @@ -37,15 +37,6 @@ def install() -> None:
__asyncio.set_event_loop_policy(EventLoopPolicy())


@_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],
*,
Expand All @@ -64,35 +55,15 @@ async def wrapper() -> _T:
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))
if sys.version_info >= (3, 12):
return __asyncio.run(
wrapper(),
loop_factory=loop_factory,
debug=debug,
**run_kwargs
)

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())
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 sys.version_info >= (3, 11):
if __asyncio._get_running_loop() is not None:
raise RuntimeError(
"asyncio.run() cannot be called from a running event loop")
Expand All @@ -104,14 +75,30 @@ async def wrapper() -> _T:
) as runner:
return runner.run(wrapper())

else:
assert vi >= (3, 12)
return __asyncio.run(
wrapper(),
loop_factory=loop_factory,
debug=debug,
**run_kwargs
)
# 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())
if hasattr(loop, 'shutdown_default_executor'):
loop.run_until_complete(loop.shutdown_default_executor())
finally:
__asyncio.set_event_loop(None)
loop.close()


def _cancel_all_tasks(loop: __asyncio.AbstractEventLoop) -> None:
Expand Down

0 comments on commit d887fa4

Please sign in to comment.