Skip to content

Commit

Permalink
Modernize event loop behavior (#387)
Browse files Browse the repository at this point in the history
* Modernize event loop behavior

* fixups

* fix downstream test

* rename method
  • Loading branch information
blink1073 authored Jan 3, 2024
1 parent bab5464 commit 19912f4
Show file tree
Hide file tree
Showing 5 changed files with 148 additions and 16 deletions.
16 changes: 16 additions & 0 deletions .github/workflows/downstream.yml
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,21 @@ jobs:
with:
package_name: jupyter_client

pytest_jupyter:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v4

- name: Base Setup
uses: jupyterlab/maintainer-tools/.github/actions/base-setup@v1

- name: Run Test
uses: jupyterlab/maintainer-tools/.github/actions/downstream-test@v1
with:
package_name: pytest_jupyter
package_spec: pip install -e ".[test,client,server]"

downstreams_check: # This job does nothing and is only used for the branch protection
if: always()
needs:
Expand All @@ -74,6 +89,7 @@ jobs:
- nbconvert
- jupyter_server
- jupyter_client
- pytest_jupyter
runs-on: ubuntu-latest
steps:
- name: Decide whether the needed jobs succeeded or failed
Expand Down
33 changes: 32 additions & 1 deletion jupyter_core/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
jupyter_path,
jupyter_runtime_dir,
)
from .utils import ensure_dir_exists
from .utils import ensure_dir_exists, ensure_event_loop

# mypy: disable-error-code="no-untyped-call"

Expand Down Expand Up @@ -277,10 +277,41 @@ def start(self) -> None:
@classmethod
def launch_instance(cls, argv: t.Any = None, **kwargs: t.Any) -> None:
"""Launch an instance of a Jupyter Application"""
# Ensure an event loop is set before any other code runs.
loop = ensure_event_loop()
try:
super().launch_instance(argv=argv, **kwargs)
except NoStart:
return
loop.close()


class JupyterAsyncApp(Application):
"""A Jupyter application that runs on an asyncio loop."""

# Set to True for tornado-based apps.
_prefer_selector_loop = False

async def initialize_async(self, argv: t.Any = None) -> None:
"""Initialize the application asynchronoously."""

async def start_async(self) -> None:
"""Run the application in an event loop."""

@classmethod
async def _launch_instance(cls, argv: t.Any = None, **kwargs: t.Any) -> None:
app = cls.instance(**kwargs)
app.initialize(argv)
await app.initialize_async(argv)
await app.start_async()

@classmethod
def launch_instance(cls, argv: t.Any = None, **kwargs: t.Any) -> None:
"""Launch an instance of an async Jupyter Application"""
loop = ensure_event_loop(cls._prefer_selector_loop)
coro = cls._launch_instance(argv, **kwargs)
loop.run_until_complete(coro)
loop.close()


if __name__ == "__main__":
Expand Down
32 changes: 21 additions & 11 deletions jupyter_core/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import sys
import threading
import warnings
from contextvars import ContextVar
from pathlib import Path
from types import FrameType
from typing import Any, Awaitable, Callable, TypeVar, cast
Expand Down Expand Up @@ -126,6 +127,7 @@ def run(self, coro: Any) -> Any:


_runner_map: dict[str, _TaskRunner] = {}
_loop: ContextVar[asyncio.AbstractEventLoop | None] = ContextVar("_loop", default=None)


def run_sync(coro: Callable[..., Awaitable[T]]) -> Callable[..., T]:
Expand Down Expand Up @@ -159,22 +161,30 @@ def wrapped(*args: Any, **kwargs: Any) -> Any:
pass

# Run the loop for this thread.
# In Python 3.12, a deprecation warning is raised, which
# may later turn into a RuntimeError. We handle both
# cases.
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
try:
loop = asyncio.get_event_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
return loop.run_until_complete(inner)
loop = ensure_event_loop()
return loop.run_until_complete(inner)

wrapped.__doc__ = coro.__doc__
return wrapped


def ensure_event_loop(prefer_selector_loop: bool = False) -> asyncio.AbstractEventLoop:
# Get the loop for this thread, or create a new one.
loop = _loop.get()
if loop is not None and not loop.is_closed():
return loop
try:
loop = asyncio.get_running_loop()
except RuntimeError:
if sys.platform == "win32" and prefer_selector_loop:
loop = asyncio.WindowsSelectorEventLoopPolicy().new_event_loop()
else:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
_loop.set(loop)
return loop


async def ensure_async(obj: Awaitable[T] | T) -> T:
"""Convert a non-awaitable object to a coroutine if needed,
and await it if it was not already awaited.
Expand Down
61 changes: 60 additions & 1 deletion tests/test_application.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import asyncio
import os
import shutil
from tempfile import mkdtemp
Expand All @@ -8,7 +9,8 @@
import pytest
from traitlets import Integer

from jupyter_core.application import JupyterApp, NoStart
from jupyter_core.application import JupyterApp, JupyterAsyncApp, NoStart
from jupyter_core.utils import ensure_event_loop

pjoin = os.path.join

Expand Down Expand Up @@ -125,3 +127,60 @@ def test_runtime_dir_changed():
app.runtime_dir = td
assert os.path.isdir(td)
shutil.rmtree(td)


class AsyncioRunApp(JupyterApp):
async def _inner(self):
pass

def start(self):
asyncio.run(self._inner())


def test_asyncio_run():
AsyncioRunApp.launch_instance([])
AsyncioRunApp.clear_instance()


class SyncTornadoApp(JupyterApp):
async def _inner(self):
self.running_loop = asyncio.get_running_loop()

def start(self):
self.starting_loop = ensure_event_loop()
loop = asyncio.get_event_loop()
loop.run_until_complete(self._inner())
loop.close()


def test_sync_tornado_run():
SyncTornadoApp.launch_instance([])
app = SyncTornadoApp.instance()
assert app.running_loop == app.starting_loop
SyncTornadoApp.clear_instance()


class AsyncApp(JupyterAsyncApp):
async def initialize_async(self, argv):
self.value = 10

async def start_async(self):
assert self.value == 10


def test_async_app():
AsyncApp.launch_instance([])
app = AsyncApp.instance()
assert app.value == 10
AsyncApp.clear_instance()


class AsyncTornadoApp(AsyncApp):
_prefer_selector_loop = True


def test_async_tornado_app():
AsyncTornadoApp.launch_instance([])
app = AsyncApp.instance()
assert app._prefer_selector_loop is True
AsyncTornadoApp.clear_instance()
22 changes: 19 additions & 3 deletions tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,13 @@

import pytest

from jupyter_core.utils import deprecation, ensure_async, ensure_dir_exists, run_sync
from jupyter_core.utils import (
deprecation,
ensure_async,
ensure_dir_exists,
ensure_event_loop,
run_sync,
)


def test_ensure_dir_exists():
Expand Down Expand Up @@ -42,11 +48,11 @@ async def foo():
foo_sync = run_sync(foo)
assert foo_sync() == 1
assert foo_sync() == 1
asyncio.get_event_loop().close()
ensure_event_loop().close()

asyncio.set_event_loop(None)
assert foo_sync() == 1
asyncio.get_event_loop().close()
ensure_event_loop().close()

asyncio.run(foo())

Expand All @@ -57,3 +63,13 @@ async def main():
assert await ensure_async(func()) == "func"

asyncio.run(main())


def test_ensure_event_loop():
loop = ensure_event_loop()

async def inner():
return asyncio.get_running_loop()

inner_sync = run_sync(inner)
assert inner_sync() == loop

0 comments on commit 19912f4

Please sign in to comment.