feat: expose Jupyter shell through web app
This commit is contained in:
@@ -0,0 +1,71 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from uuid import uuid4
|
||||
|
||||
from .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
|
||||
Reference in New Issue
Block a user