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."""
|
"""Small, side-effect-free API for the in-process IPython runner."""
|
||||||
|
|
||||||
from .app import App
|
from .app import App
|
||||||
|
from .jupyter_app import JupyterApp
|
||||||
|
from .jupyter_shell import JupyterShell
|
||||||
from .events import (
|
from .events import (
|
||||||
error_from_execution_result,
|
error_from_execution_result,
|
||||||
event_from_execution_result,
|
event_from_execution_result,
|
||||||
@@ -26,10 +28,16 @@ from .shell import (
|
|||||||
run_cell_and_collect,
|
run_cell_and_collect,
|
||||||
setup_shell,
|
setup_shell,
|
||||||
)
|
)
|
||||||
|
from .transport import InputRequest, JupyterTransport, KernelMessage
|
||||||
from .utils import generate_good_names
|
from .utils import generate_good_names
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"App",
|
"App",
|
||||||
|
"JupyterApp",
|
||||||
|
"JupyterShell",
|
||||||
|
"JupyterTransport",
|
||||||
|
"InputRequest",
|
||||||
|
"KernelMessage",
|
||||||
"AppInfo",
|
"AppInfo",
|
||||||
"CallError",
|
"CallError",
|
||||||
"CallEvent",
|
"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