Skip to content

Commit 19912f4

Browse files
authored
Modernize event loop behavior (#387)
* Modernize event loop behavior * fixups * fix downstream test * rename method
1 parent bab5464 commit 19912f4

File tree

5 files changed

+148
-16
lines changed

5 files changed

+148
-16
lines changed

.github/workflows/downstream.yml

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,21 @@ jobs:
6666
with:
6767
package_name: jupyter_client
6868

69+
pytest_jupyter:
70+
runs-on: ubuntu-latest
71+
steps:
72+
- name: Checkout
73+
uses: actions/checkout@v4
74+
75+
- name: Base Setup
76+
uses: jupyterlab/maintainer-tools/.github/actions/base-setup@v1
77+
78+
- name: Run Test
79+
uses: jupyterlab/maintainer-tools/.github/actions/downstream-test@v1
80+
with:
81+
package_name: pytest_jupyter
82+
package_spec: pip install -e ".[test,client,server]"
83+
6984
downstreams_check: # This job does nothing and is only used for the branch protection
7085
if: always()
7186
needs:
@@ -74,6 +89,7 @@ jobs:
7489
- nbconvert
7590
- jupyter_server
7691
- jupyter_client
92+
- pytest_jupyter
7793
runs-on: ubuntu-latest
7894
steps:
7995
- name: Decide whether the needed jobs succeeded or failed

jupyter_core/application.py

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
jupyter_path,
3030
jupyter_runtime_dir,
3131
)
32-
from .utils import ensure_dir_exists
32+
from .utils import ensure_dir_exists, ensure_event_loop
3333

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

@@ -277,10 +277,41 @@ def start(self) -> None:
277277
@classmethod
278278
def launch_instance(cls, argv: t.Any = None, **kwargs: t.Any) -> None:
279279
"""Launch an instance of a Jupyter Application"""
280+
# Ensure an event loop is set before any other code runs.
281+
loop = ensure_event_loop()
280282
try:
281283
super().launch_instance(argv=argv, **kwargs)
282284
except NoStart:
283285
return
286+
loop.close()
287+
288+
289+
class JupyterAsyncApp(Application):
290+
"""A Jupyter application that runs on an asyncio loop."""
291+
292+
# Set to True for tornado-based apps.
293+
_prefer_selector_loop = False
294+
295+
async def initialize_async(self, argv: t.Any = None) -> None:
296+
"""Initialize the application asynchronoously."""
297+
298+
async def start_async(self) -> None:
299+
"""Run the application in an event loop."""
300+
301+
@classmethod
302+
async def _launch_instance(cls, argv: t.Any = None, **kwargs: t.Any) -> None:
303+
app = cls.instance(**kwargs)
304+
app.initialize(argv)
305+
await app.initialize_async(argv)
306+
await app.start_async()
307+
308+
@classmethod
309+
def launch_instance(cls, argv: t.Any = None, **kwargs: t.Any) -> None:
310+
"""Launch an instance of an async Jupyter Application"""
311+
loop = ensure_event_loop(cls._prefer_selector_loop)
312+
coro = cls._launch_instance(argv, **kwargs)
313+
loop.run_until_complete(coro)
314+
loop.close()
284315

285316

286317
if __name__ == "__main__":

jupyter_core/utils/__init__.py

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import sys
1010
import threading
1111
import warnings
12+
from contextvars import ContextVar
1213
from pathlib import Path
1314
from types import FrameType
1415
from typing import Any, Awaitable, Callable, TypeVar, cast
@@ -126,6 +127,7 @@ def run(self, coro: Any) -> Any:
126127

127128

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

130132

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

161163
# Run the loop for this thread.
162-
# In Python 3.12, a deprecation warning is raised, which
163-
# may later turn into a RuntimeError. We handle both
164-
# cases.
165-
with warnings.catch_warnings():
166-
warnings.simplefilter("ignore", DeprecationWarning)
167-
try:
168-
loop = asyncio.get_event_loop()
169-
except RuntimeError:
170-
loop = asyncio.new_event_loop()
171-
asyncio.set_event_loop(loop)
172-
return loop.run_until_complete(inner)
164+
loop = ensure_event_loop()
165+
return loop.run_until_complete(inner)
173166

174167
wrapped.__doc__ = coro.__doc__
175168
return wrapped
176169

177170

171+
def ensure_event_loop(prefer_selector_loop: bool = False) -> asyncio.AbstractEventLoop:
172+
# Get the loop for this thread, or create a new one.
173+
loop = _loop.get()
174+
if loop is not None and not loop.is_closed():
175+
return loop
176+
try:
177+
loop = asyncio.get_running_loop()
178+
except RuntimeError:
179+
if sys.platform == "win32" and prefer_selector_loop:
180+
loop = asyncio.WindowsSelectorEventLoopPolicy().new_event_loop()
181+
else:
182+
loop = asyncio.new_event_loop()
183+
asyncio.set_event_loop(loop)
184+
_loop.set(loop)
185+
return loop
186+
187+
178188
async def ensure_async(obj: Awaitable[T] | T) -> T:
179189
"""Convert a non-awaitable object to a coroutine if needed,
180190
and await it if it was not already awaited.

tests/test_application.py

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import asyncio
34
import os
45
import shutil
56
from tempfile import mkdtemp
@@ -8,7 +9,8 @@
89
import pytest
910
from traitlets import Integer
1011

11-
from jupyter_core.application import JupyterApp, NoStart
12+
from jupyter_core.application import JupyterApp, JupyterAsyncApp, NoStart
13+
from jupyter_core.utils import ensure_event_loop
1214

1315
pjoin = os.path.join
1416

@@ -125,3 +127,60 @@ def test_runtime_dir_changed():
125127
app.runtime_dir = td
126128
assert os.path.isdir(td)
127129
shutil.rmtree(td)
130+
131+
132+
class AsyncioRunApp(JupyterApp):
133+
async def _inner(self):
134+
pass
135+
136+
def start(self):
137+
asyncio.run(self._inner())
138+
139+
140+
def test_asyncio_run():
141+
AsyncioRunApp.launch_instance([])
142+
AsyncioRunApp.clear_instance()
143+
144+
145+
class SyncTornadoApp(JupyterApp):
146+
async def _inner(self):
147+
self.running_loop = asyncio.get_running_loop()
148+
149+
def start(self):
150+
self.starting_loop = ensure_event_loop()
151+
loop = asyncio.get_event_loop()
152+
loop.run_until_complete(self._inner())
153+
loop.close()
154+
155+
156+
def test_sync_tornado_run():
157+
SyncTornadoApp.launch_instance([])
158+
app = SyncTornadoApp.instance()
159+
assert app.running_loop == app.starting_loop
160+
SyncTornadoApp.clear_instance()
161+
162+
163+
class AsyncApp(JupyterAsyncApp):
164+
async def initialize_async(self, argv):
165+
self.value = 10
166+
167+
async def start_async(self):
168+
assert self.value == 10
169+
170+
171+
def test_async_app():
172+
AsyncApp.launch_instance([])
173+
app = AsyncApp.instance()
174+
assert app.value == 10
175+
AsyncApp.clear_instance()
176+
177+
178+
class AsyncTornadoApp(AsyncApp):
179+
_prefer_selector_loop = True
180+
181+
182+
def test_async_tornado_app():
183+
AsyncTornadoApp.launch_instance([])
184+
app = AsyncApp.instance()
185+
assert app._prefer_selector_loop is True
186+
AsyncTornadoApp.clear_instance()

tests/test_utils.py

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,13 @@
1010

1111
import pytest
1212

13-
from jupyter_core.utils import deprecation, ensure_async, ensure_dir_exists, run_sync
13+
from jupyter_core.utils import (
14+
deprecation,
15+
ensure_async,
16+
ensure_dir_exists,
17+
ensure_event_loop,
18+
run_sync,
19+
)
1420

1521

1622
def test_ensure_dir_exists():
@@ -42,11 +48,11 @@ async def foo():
4248
foo_sync = run_sync(foo)
4349
assert foo_sync() == 1
4450
assert foo_sync() == 1
45-
asyncio.get_event_loop().close()
51+
ensure_event_loop().close()
4652

4753
asyncio.set_event_loop(None)
4854
assert foo_sync() == 1
49-
asyncio.get_event_loop().close()
55+
ensure_event_loop().close()
5056

5157
asyncio.run(foo())
5258

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

5965
asyncio.run(main())
66+
67+
68+
def test_ensure_event_loop():
69+
loop = ensure_event_loop()
70+
71+
async def inner():
72+
return asyncio.get_running_loop()
73+
74+
inner_sync = run_sync(inner)
75+
assert inner_sync() == loop

0 commit comments

Comments
 (0)