diff --git a/src/ipython_shell/__init__.py b/src/ipython_shell/__init__.py index 9137e59..8cf4763 100644 --- a/src/ipython_shell/__init__.py +++ b/src/ipython_shell/__init__.py @@ -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", diff --git a/src/ipython_shell/jupyter_app.py b/src/ipython_shell/jupyter_app.py new file mode 100644 index 0000000..35e3caf --- /dev/null +++ b/src/ipython_shell/jupyter_app.py @@ -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 diff --git a/src/ipython_shell/jupyter_shell.py b/src/ipython_shell/jupyter_shell.py new file mode 100644 index 0000000..ac2a100 --- /dev/null +++ b/src/ipython_shell/jupyter_shell.py @@ -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() diff --git a/tests/test_jupyter_app.py b/tests/test_jupyter_app.py new file mode 100644 index 0000000..f54f2d9 --- /dev/null +++ b/tests/test_jupyter_app.py @@ -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()