diff --git a/src/ipython_shell/transport.py b/src/ipython_shell/transport.py index b1106ad..3745479 100644 --- a/src/ipython_shell/transport.py +++ b/src/ipython_shell/transport.py @@ -32,6 +32,7 @@ class JupyterTransport: self.client = None self._active_call: str | None = None self._call_lock = asyncio.Lock() + self._waiting_for_input: str | None = None async def start(self) -> None: """Start the kernel process and wait until it accepts requests.""" @@ -51,8 +52,10 @@ class JupyterTransport: self._active_call = self.client.execute(code, allow_stdin=True) return self._active_call - async def messages_for(self, call_id: str) -> AsyncIterator[KernelMessage]: - """Yield decoded IOPub and shell messages for one execution.""" + async def messages_for( + self, call_id: str + ) -> AsyncIterator[KernelMessage | InputRequest]: + """Yield decoded output, input, and completion messages for one call.""" if self.client is None: raise RuntimeError("JupyterTransport has not been started") if self._active_call != call_id: @@ -64,8 +67,9 @@ class JupyterTransport: while True: iopub_task = asyncio.create_task(self.client.get_iopub_msg()) shell_task = asyncio.create_task(self.client.get_shell_msg()) + stdin_task = asyncio.create_task(self.client.get_stdin_msg()) done, pending = await asyncio.wait( - {iopub_task, shell_task}, + {iopub_task, shell_task, stdin_task}, return_when=asyncio.FIRST_COMPLETED, ) for task in pending: @@ -78,6 +82,17 @@ class JupyterTransport: if parent_id != call_id: continue + if message["msg_type"] == "input_request": + self._waiting_for_input = call_id + yield InputRequest( + call_id=call_id, + prompt=str(message.get("content", {}).get("prompt", "")), + password=bool( + message.get("content", {}).get("password", False) + ), + ) + continue + decoded = KernelMessage( msg_type=message["msg_type"], parent_id=parent_id, @@ -101,9 +116,21 @@ class JupyterTransport: yield reply return finally: + if self._waiting_for_input == call_id: + self._waiting_for_input = None if self._active_call == call_id: self._active_call = None + async def reply_to_input(self, value: str) -> None: + """Reply to the active kernel input request without using process stdin.""" + if self.client is None: + raise RuntimeError("JupyterTransport has not been started") + if self._waiting_for_input is None: + raise RuntimeError("JupyterTransport is not waiting for input") + + self.client.input(value) + self._waiting_for_input = None + async def shutdown(self) -> None: """Stop channels and terminate the kernel process.""" if self.client is not None: diff --git a/tests/test_transport.py b/tests/test_transport.py index 4d5c8b4..78d1b23 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -1,6 +1,7 @@ +import asyncio import unittest -from ipython_shell.transport import JupyterTransport, KernelMessage +from ipython_shell.transport import InputRequest, JupyterTransport, KernelMessage class TransportTests(unittest.TestCase): @@ -35,6 +36,60 @@ class AsyncTransportTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(messages[-1].msg_type, "execute_reply") self.assertEqual(messages[-1].content["status"], "ok") + async def test_transport_preserves_mime_and_errors(self): + transport = JupyterTransport() + await transport.start() + try: + display_id = await transport.execute( + "from IPython.display import display\n" + "display({'text/plain': 'hello'}, raw=True)" + ) + display_messages = [ + message async for message in transport.messages_for(display_id) + ] + + error_id = await transport.execute("raise ValueError('boom')") + error_messages = [ + message async for message in transport.messages_for(error_id) + ] + finally: + await transport.shutdown() + + display = next( + message + for message in display_messages + if message.msg_type == "display_data" + ) + self.assertEqual(display.content["data"]["text/plain"], "hello") + + error = next( + message for message in error_messages if message.msg_type == "error" + ) + self.assertEqual(error.content["ename"], "ValueError") + + async def test_transport_routes_input_reply_without_parent_stdin(self): + transport = JupyterTransport() + await transport.start() + try: + call_id = await transport.execute("answer = input('name? '); answer") + stream = transport.messages_for(call_id) + + while True: + event = await asyncio.wait_for(anext(stream), timeout=5) + if isinstance(event, InputRequest): + input_request = event + break + + self.assertEqual(input_request.prompt, "name? ") + await transport.reply_to_input("Ada") + messages = [message async for message in stream] + finally: + await transport.shutdown() + + self.assertTrue( + any(message.msg_type == "execute_result" for message in messages) + ) + if __name__ == "__main__": unittest.main()