feat: add Jupyter-backed shell app
This commit is contained in:
@@ -1,6 +1,8 @@
|
||||
"""Small, side-effect-free API for the in-process IPython runner."""
|
||||
|
||||
from .app import App
|
||||
from .jupyter_app import JupyterApp
|
||||
from .jupyter_shell import JupyterShell
|
||||
from .events import (
|
||||
error_from_execution_result,
|
||||
event_from_execution_result,
|
||||
@@ -26,10 +28,16 @@ from .shell import (
|
||||
run_cell_and_collect,
|
||||
setup_shell,
|
||||
)
|
||||
from .transport import InputRequest, JupyterTransport, KernelMessage
|
||||
from .utils import generate_good_names
|
||||
|
||||
__all__ = [
|
||||
"App",
|
||||
"JupyterApp",
|
||||
"JupyterShell",
|
||||
"JupyterTransport",
|
||||
"InputRequest",
|
||||
"KernelMessage",
|
||||
"AppInfo",
|
||||
"CallError",
|
||||
"CallEvent",
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from uuid import uuid4
|
||||
|
||||
from .jupyter_shell import JupyterShell
|
||||
from .transport import InputRequest, KernelMessage
|
||||
from .utils import generate_good_names
|
||||
|
||||
|
||||
class JupyterApp:
|
||||
"""Coordinate named persistent shells backed by Jupyter kernels."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.shells: dict[str, JupyterShell] = {}
|
||||
self.last_shell: str | None = None
|
||||
self._calls: dict[str, JupyterShell] = {}
|
||||
|
||||
def _new_shell(self) -> tuple[str, JupyterShell]:
|
||||
name = generate_good_names()
|
||||
while name in self.shells:
|
||||
name = generate_good_names()
|
||||
shell = JupyterShell(name)
|
||||
self.shells[name] = shell
|
||||
self.last_shell = name
|
||||
return name, shell
|
||||
|
||||
def _select_shell(self, shell_name: str) -> tuple[str, JupyterShell]:
|
||||
if shell_name == "new":
|
||||
return self._new_shell()
|
||||
if shell_name == "last":
|
||||
if self.last_shell is None:
|
||||
return self._new_shell()
|
||||
return self.last_shell, self.shells[self.last_shell]
|
||||
if shell_name not in self.shells:
|
||||
self.shells[shell_name] = JupyterShell(shell_name)
|
||||
self.last_shell = shell_name
|
||||
return shell_name, self.shells[shell_name]
|
||||
|
||||
def get_shell(self, shell_name: str) -> JupyterShell:
|
||||
"""Return a named Jupyter shell."""
|
||||
if shell_name not in self.shells:
|
||||
raise KeyError(f"Shell {shell_name} does not exist")
|
||||
return self.shells[shell_name]
|
||||
|
||||
async def run_code_stream(
|
||||
self,
|
||||
code: str,
|
||||
shell_name: str = "last",
|
||||
) -> AsyncIterator[KernelMessage | InputRequest]:
|
||||
"""Execute code and yield its subprocess messages."""
|
||||
_, shell = self._select_shell(shell_name)
|
||||
call_id = uuid4().hex
|
||||
self._calls[call_id] = shell
|
||||
try:
|
||||
async for message in shell.run_cell_stream(code, call_id=call_id):
|
||||
yield message
|
||||
finally:
|
||||
self._calls.pop(call_id, None)
|
||||
|
||||
async def reply_to_input(self, call_id: str, value: str) -> None:
|
||||
"""Reply to an input request emitted by a running call."""
|
||||
try:
|
||||
shell = self._calls[call_id]
|
||||
except KeyError as error:
|
||||
raise KeyError(f"Call {call_id} does not exist") from error
|
||||
await shell.reply_to_input(call_id, value)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Stop all kernel processes owned by this app."""
|
||||
awaitables = [shell.shutdown() for shell in self.shells.values()]
|
||||
for awaitable in awaitables:
|
||||
await awaitable
|
||||
@@ -0,0 +1,49 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import replace
|
||||
|
||||
from .transport import InputRequest, JupyterTransport, KernelMessage
|
||||
|
||||
|
||||
class JupyterShell:
|
||||
"""One persistent, subprocess-backed IPython shell."""
|
||||
|
||||
def __init__(self, name: str, transport: JupyterTransport | None = None) -> None:
|
||||
self.name = name
|
||||
self.transport = transport or JupyterTransport()
|
||||
self._active_call_id: str | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start this shell's kernel if it is not already running."""
|
||||
if self.transport.client is None:
|
||||
await self.transport.start()
|
||||
|
||||
async def run_cell_stream(
|
||||
self,
|
||||
code: str,
|
||||
*,
|
||||
call_id: str,
|
||||
) -> AsyncIterator[KernelMessage | InputRequest]:
|
||||
"""Yield kernel messages while one cell executes."""
|
||||
await self.start()
|
||||
self._active_call_id = call_id
|
||||
try:
|
||||
kernel_call_id = await self.transport.execute(code)
|
||||
async for message in self.transport.messages_for(kernel_call_id):
|
||||
if isinstance(message, InputRequest):
|
||||
yield replace(message, call_id=call_id)
|
||||
else:
|
||||
# Keep the kernel's execution ID private to the transport.
|
||||
yield replace(message, parent_id=call_id)
|
||||
finally:
|
||||
if self._active_call_id == call_id:
|
||||
self._active_call_id = None
|
||||
|
||||
async def reply_to_input(self, call_id: str, value: str) -> None:
|
||||
"""Send input to the cell currently waiting in this shell."""
|
||||
if self._active_call_id != call_id:
|
||||
raise RuntimeError(f"No active input request for call: {call_id}")
|
||||
await self.transport.reply_to_input(value)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Stop this shell's kernel process."""
|
||||
await self.transport.shutdown()
|
||||
@@ -0,0 +1,87 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
|
||||
from ipython_shell.jupyter_app import JupyterApp
|
||||
from ipython_shell.jupyter_shell import JupyterShell
|
||||
from ipython_shell.transport import InputRequest, KernelMessage
|
||||
|
||||
|
||||
class JupyterShellTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_shell_keeps_namespace_between_cells(self):
|
||||
shell = JupyterShell("persistent")
|
||||
try:
|
||||
first = [
|
||||
message
|
||||
async for message in shell.run_cell_stream(
|
||||
"answer = 41; answer",
|
||||
call_id="call-1",
|
||||
)
|
||||
]
|
||||
second = [
|
||||
message
|
||||
async for message in shell.run_cell_stream(
|
||||
"answer + 1",
|
||||
call_id="call-2",
|
||||
)
|
||||
]
|
||||
finally:
|
||||
await shell.shutdown()
|
||||
|
||||
self.assertEqual(
|
||||
next(
|
||||
message.content["data"]["text/plain"]
|
||||
for message in first
|
||||
if isinstance(message, KernelMessage)
|
||||
and message.msg_type == "execute_result"
|
||||
),
|
||||
"41",
|
||||
)
|
||||
self.assertEqual(
|
||||
next(
|
||||
message.content["data"]["text/plain"]
|
||||
for message in second
|
||||
if isinstance(message, KernelMessage)
|
||||
and message.msg_type == "execute_result"
|
||||
),
|
||||
"42",
|
||||
)
|
||||
|
||||
|
||||
class JupyterAppTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_app_routes_input_reply_to_the_call(self):
|
||||
app = JupyterApp()
|
||||
stream = app.run_code_stream(
|
||||
"answer = input('name? '); answer",
|
||||
shell_name="new",
|
||||
)
|
||||
|
||||
try:
|
||||
while True:
|
||||
event = await asyncio.wait_for(anext(stream), timeout=5)
|
||||
if isinstance(event, InputRequest):
|
||||
input_request = event
|
||||
break
|
||||
|
||||
self.assertIsInstance(input_request, InputRequest)
|
||||
assert isinstance(input_request, InputRequest)
|
||||
self.assertTrue(input_request.call_id)
|
||||
self.assertEqual(input_request.prompt, "name? ")
|
||||
|
||||
await app.reply_to_input(input_request.call_id, "Ada")
|
||||
messages = [message async for message in stream]
|
||||
finally:
|
||||
await app.shutdown()
|
||||
|
||||
self.assertEqual(
|
||||
next(
|
||||
message.content["data"]["text/plain"]
|
||||
for message in messages
|
||||
if isinstance(message, KernelMessage)
|
||||
and message.msg_type == "execute_result"
|
||||
),
|
||||
"'Ada'",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user