egutierrez
a9f2c60e3d
jupyter_discover: soporte multi-servidor, detección de modo colaborativo mejorada. jupyter_write: operaciones batch (insert, edit, delete), manejo robusto de Y.js. jupyter_exec: mejoras en ejecución directa al kernel. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
504 lines
16 KiB
Python
504 lines
16 KiB
Python
"""Operaciones de escritura sobre celdas de un notebook Jupyter via colaboracion en tiempo real.
|
|
|
|
NO ejecuta celdas — solo modifica la estructura del notebook (append, insert, edit, delete, create, batch).
|
|
Usa jupyter_nbmodel_client para comunicarse con el servidor Jupyter via WebSocket.
|
|
Para crear notebooks usa la API REST de Jupyter (urllib PUT /api/contents).
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import argparse
|
|
import sys
|
|
from urllib.error import URLError, HTTPError
|
|
from urllib.request import Request, urlopen
|
|
|
|
from jupyter_nbmodel_client import NbModelClient, get_jupyter_notebook_websocket_url
|
|
|
|
|
|
def _resolve_collab_username(server_url: str, token: str) -> str:
|
|
"""Resolve the display name of the active user in Jupyter collaboration."""
|
|
headers = {"Accept": "application/json"}
|
|
if token:
|
|
headers["Authorization"] = f"token {token}"
|
|
try:
|
|
req = Request(f"{server_url}/api/me", headers=headers)
|
|
with urlopen(req, timeout=5) as resp:
|
|
me = json.loads(resp.read())
|
|
identity = me.get("identity", {})
|
|
return identity.get("display_name", "") or identity.get("username", "") or identity.get("name", "Anonymous")
|
|
except (URLError, OSError, json.JSONDecodeError):
|
|
return "Anonymous"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers internos async
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _append_cell(
|
|
notebook_path: str,
|
|
source: str,
|
|
cell_type: str,
|
|
server_url: str,
|
|
token: str,
|
|
) -> dict:
|
|
ws_url = get_jupyter_notebook_websocket_url(
|
|
server_url=server_url,
|
|
token=token,
|
|
path=notebook_path,
|
|
)
|
|
username = _resolve_collab_username(server_url, token)
|
|
async with NbModelClient(ws_url, username=username) as nb:
|
|
if cell_type == "markdown":
|
|
nb.add_markdown_cell(source)
|
|
else:
|
|
nb.add_code_cell(source)
|
|
cell_index = len(nb) - 1
|
|
await asyncio.sleep(2)
|
|
return {
|
|
"action": f"append_{cell_type}",
|
|
"cell_index": cell_index,
|
|
"notebook": notebook_path,
|
|
}
|
|
|
|
|
|
async def _insert_cell(
|
|
notebook_path: str,
|
|
cell_index: int,
|
|
source: str,
|
|
cell_type: str,
|
|
server_url: str,
|
|
token: str,
|
|
) -> dict:
|
|
ws_url = get_jupyter_notebook_websocket_url(
|
|
server_url=server_url,
|
|
token=token,
|
|
path=notebook_path,
|
|
)
|
|
username = _resolve_collab_username(server_url, token)
|
|
async with NbModelClient(ws_url, username=username) as nb:
|
|
nb.insert_cell(cell_index, cell_type=cell_type, source=source)
|
|
await asyncio.sleep(2)
|
|
return {
|
|
"action": "insert",
|
|
"cell_index": cell_index,
|
|
"cell_type": cell_type,
|
|
"notebook": notebook_path,
|
|
}
|
|
|
|
|
|
async def _edit_cell(
|
|
notebook_path: str,
|
|
cell_index: int,
|
|
source: str,
|
|
server_url: str,
|
|
token: str,
|
|
) -> dict:
|
|
ws_url = get_jupyter_notebook_websocket_url(
|
|
server_url=server_url,
|
|
token=token,
|
|
path=notebook_path,
|
|
)
|
|
username = _resolve_collab_username(server_url, token)
|
|
async with NbModelClient(ws_url, username=username) as nb:
|
|
nb.set_cell_source(cell_index, source)
|
|
await asyncio.sleep(2)
|
|
return {
|
|
"action": "edit",
|
|
"cell_index": cell_index,
|
|
"notebook": notebook_path,
|
|
}
|
|
|
|
|
|
async def _delete_cell(
|
|
notebook_path: str,
|
|
cell_index: int,
|
|
server_url: str,
|
|
token: str,
|
|
) -> dict:
|
|
ws_url = get_jupyter_notebook_websocket_url(
|
|
server_url=server_url,
|
|
token=token,
|
|
path=notebook_path,
|
|
)
|
|
username = _resolve_collab_username(server_url, token)
|
|
async with NbModelClient(ws_url, username=username) as nb:
|
|
nb.delete_cell(cell_index)
|
|
await asyncio.sleep(2)
|
|
return {
|
|
"action": "delete",
|
|
"cell_index": cell_index,
|
|
"notebook": notebook_path,
|
|
}
|
|
|
|
|
|
async def _batch_write(
|
|
notebook_path: str,
|
|
cells: list,
|
|
server_url: str,
|
|
token: str,
|
|
) -> dict:
|
|
"""Anade multiples celdas en una sola conexion WebSocket."""
|
|
ws_url = get_jupyter_notebook_websocket_url(
|
|
server_url=server_url,
|
|
token=token,
|
|
path=notebook_path,
|
|
)
|
|
username = _resolve_collab_username(server_url, token)
|
|
async with NbModelClient(ws_url, username=username) as nb:
|
|
for cell in cells:
|
|
cell_type = cell.get("type", "code")
|
|
source = cell.get("source", "")
|
|
if cell_type == "markdown":
|
|
nb.add_markdown_cell(source)
|
|
else:
|
|
nb.add_code_cell(source)
|
|
await asyncio.sleep(2)
|
|
return {
|
|
"action": "batch",
|
|
"cells_added": len(cells),
|
|
"notebook": notebook_path,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API publica sincrona
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def jupyter_append_code(
|
|
notebook_path: str,
|
|
source: str,
|
|
server_url: str = "http://localhost:8888",
|
|
token: str = "",
|
|
) -> dict:
|
|
"""Anade una celda de codigo al final del notebook.
|
|
|
|
Args:
|
|
notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
|
|
source: Codigo fuente de la celda.
|
|
server_url: URL base del servidor Jupyter.
|
|
token: Token de autenticacion del servidor Jupyter.
|
|
|
|
Returns:
|
|
dict con action, cell_index y notebook.
|
|
"""
|
|
return asyncio.run(_append_cell(notebook_path, source, "code", server_url, token))
|
|
|
|
|
|
def jupyter_append_markdown(
|
|
notebook_path: str,
|
|
source: str,
|
|
server_url: str = "http://localhost:8888",
|
|
token: str = "",
|
|
) -> dict:
|
|
"""Anade una celda markdown al final del notebook.
|
|
|
|
Args:
|
|
notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
|
|
source: Contenido markdown de la celda.
|
|
server_url: URL base del servidor Jupyter.
|
|
token: Token de autenticacion del servidor Jupyter.
|
|
|
|
Returns:
|
|
dict con action, cell_index y notebook.
|
|
"""
|
|
return asyncio.run(
|
|
_append_cell(notebook_path, source, "markdown", server_url, token)
|
|
)
|
|
|
|
|
|
def jupyter_insert_cell(
|
|
notebook_path: str,
|
|
cell_index: int,
|
|
source: str,
|
|
cell_type: str = "code",
|
|
server_url: str = "http://localhost:8888",
|
|
token: str = "",
|
|
) -> dict:
|
|
"""Inserta una celda en una posicion especifica del notebook.
|
|
|
|
Args:
|
|
notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
|
|
cell_index: Indice donde insertar (0 = primer posicion).
|
|
source: Contenido de la celda.
|
|
cell_type: Tipo de celda: "code" o "markdown".
|
|
server_url: URL base del servidor Jupyter.
|
|
token: Token de autenticacion del servidor Jupyter.
|
|
|
|
Returns:
|
|
dict con action, cell_index, cell_type y notebook.
|
|
"""
|
|
return asyncio.run(
|
|
_insert_cell(notebook_path, cell_index, source, cell_type, server_url, token)
|
|
)
|
|
|
|
|
|
def jupyter_edit_cell(
|
|
notebook_path: str,
|
|
cell_index: int,
|
|
source: str,
|
|
server_url: str = "http://localhost:8888",
|
|
token: str = "",
|
|
) -> dict:
|
|
"""Sobrescribe el contenido de una celda existente.
|
|
|
|
Args:
|
|
notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
|
|
cell_index: Indice de la celda a editar (0-based).
|
|
source: Nuevo contenido de la celda.
|
|
server_url: URL base del servidor Jupyter.
|
|
token: Token de autenticacion del servidor Jupyter.
|
|
|
|
Returns:
|
|
dict con action, cell_index y notebook.
|
|
"""
|
|
return asyncio.run(
|
|
_edit_cell(notebook_path, cell_index, source, server_url, token)
|
|
)
|
|
|
|
|
|
def jupyter_delete_cell(
|
|
notebook_path: str,
|
|
cell_index: int,
|
|
server_url: str = "http://localhost:8888",
|
|
token: str = "",
|
|
) -> dict:
|
|
"""Elimina una celda del notebook.
|
|
|
|
Args:
|
|
notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
|
|
cell_index: Indice de la celda a eliminar (0-based).
|
|
server_url: URL base del servidor Jupyter.
|
|
token: Token de autenticacion del servidor Jupyter.
|
|
|
|
Returns:
|
|
dict con action, cell_index y notebook.
|
|
"""
|
|
return asyncio.run(
|
|
_delete_cell(notebook_path, cell_index, server_url, token)
|
|
)
|
|
|
|
|
|
def jupyter_create_notebook(
|
|
notebook_path: str,
|
|
kernel_name: str = "python3",
|
|
server_url: str = "http://localhost:8888",
|
|
token: str = "",
|
|
force: bool = False,
|
|
) -> dict:
|
|
"""Crea un notebook vacio (nbformat 4) via la API REST de Jupyter.
|
|
|
|
Usa PUT /api/contents/{path} con type: notebook. Si el notebook ya existe
|
|
y force=False, lanza un error en vez de sobreescribirlo.
|
|
|
|
Args:
|
|
notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
|
|
kernel_name: Nombre del kernel a asociar (default: python3).
|
|
server_url: URL base del servidor Jupyter.
|
|
token: Token de autenticacion del servidor Jupyter.
|
|
force: Si True, sobreescribe el notebook si ya existe.
|
|
|
|
Returns:
|
|
dict con action, notebook y created.
|
|
"""
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
}
|
|
if token:
|
|
headers["Authorization"] = f"token {token}"
|
|
|
|
# Verificar si ya existe (HEAD request)
|
|
check_url = f"{server_url}/api/contents/{notebook_path}"
|
|
check_req = Request(check_url, headers=headers, method="HEAD")
|
|
already_exists = False
|
|
try:
|
|
with urlopen(check_req, timeout=5):
|
|
already_exists = True
|
|
except HTTPError as e:
|
|
if e.code != 404:
|
|
raise
|
|
except URLError:
|
|
raise
|
|
|
|
if already_exists and not force:
|
|
raise FileExistsError(
|
|
f"Notebook ya existe: {notebook_path}. Usa force=True para sobreescribir."
|
|
)
|
|
|
|
kernel_display = {
|
|
"python3": "Python 3 (ipykernel)",
|
|
"python": "Python 3",
|
|
}.get(kernel_name, kernel_name)
|
|
|
|
notebook_content = {
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5,
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"name": kernel_name,
|
|
"display_name": kernel_display,
|
|
"language": "python",
|
|
},
|
|
"language_info": {
|
|
"name": "python",
|
|
},
|
|
},
|
|
"cells": [],
|
|
}
|
|
|
|
body = json.dumps({
|
|
"type": "notebook",
|
|
"content": notebook_content,
|
|
}).encode("utf-8")
|
|
|
|
put_req = Request(check_url, data=body, headers=headers, method="PUT")
|
|
with urlopen(put_req, timeout=10) as resp:
|
|
resp.read()
|
|
|
|
return {
|
|
"action": "create",
|
|
"notebook": notebook_path,
|
|
"created": True,
|
|
}
|
|
|
|
|
|
def jupyter_batch_write(
|
|
notebook_path: str,
|
|
cells: list,
|
|
server_url: str = "http://localhost:8888",
|
|
token: str = "",
|
|
) -> dict:
|
|
"""Anade multiples celdas al notebook en una sola conexion WebSocket.
|
|
|
|
Mucho mas eficiente que llamar append-code/append-markdown N veces porque
|
|
abre una unica conexion WebSocket y hace un solo sleep de sincronizacion
|
|
al final de todas las inserciones.
|
|
|
|
Args:
|
|
notebook_path: Ruta relativa al notebook dentro del servidor Jupyter.
|
|
cells: Lista de dicts con claves "type" ("code"|"markdown") y "source" (str).
|
|
server_url: URL base del servidor Jupyter.
|
|
token: Token de autenticacion del servidor Jupyter.
|
|
|
|
Returns:
|
|
dict con action, cells_added y notebook.
|
|
"""
|
|
return asyncio.run(_batch_write(notebook_path, cells, server_url, token))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
prog="jupyter_write",
|
|
description="Operaciones de escritura sobre celdas de un notebook Jupyter.",
|
|
)
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
|
|
# Argumentos comunes
|
|
def add_common(p: argparse.ArgumentParser) -> None:
|
|
p.add_argument("--server", default="http://localhost:8888", help="URL del servidor Jupyter")
|
|
p.add_argument("--token", default="", help="Token de autenticacion")
|
|
|
|
# append-code
|
|
p_ac = sub.add_parser("append-code", help="Anade celda de codigo al final")
|
|
p_ac.add_argument("notebook", help="Ruta del notebook")
|
|
p_ac.add_argument("source", help="Codigo fuente")
|
|
add_common(p_ac)
|
|
|
|
# append-markdown
|
|
p_am = sub.add_parser("append-markdown", help="Anade celda markdown al final")
|
|
p_am.add_argument("notebook", help="Ruta del notebook")
|
|
p_am.add_argument("source", help="Contenido markdown")
|
|
add_common(p_am)
|
|
|
|
# insert
|
|
p_ins = sub.add_parser("insert", help="Inserta celda en posicion especifica")
|
|
p_ins.add_argument("notebook", help="Ruta del notebook")
|
|
p_ins.add_argument("index", type=int, help="Indice de insercion (0-based)")
|
|
p_ins.add_argument("source", help="Contenido de la celda")
|
|
p_ins.add_argument("--type", dest="cell_type", choices=["code", "markdown"], default="code")
|
|
add_common(p_ins)
|
|
|
|
# edit
|
|
p_ed = sub.add_parser("edit", help="Sobrescribe el contenido de una celda")
|
|
p_ed.add_argument("notebook", help="Ruta del notebook")
|
|
p_ed.add_argument("index", type=int, help="Indice de la celda (0-based)")
|
|
p_ed.add_argument("source", help="Nuevo contenido")
|
|
add_common(p_ed)
|
|
|
|
# delete
|
|
p_del = sub.add_parser("delete", help="Elimina una celda")
|
|
p_del.add_argument("notebook", help="Ruta del notebook")
|
|
p_del.add_argument("index", type=int, help="Indice de la celda (0-based)")
|
|
add_common(p_del)
|
|
|
|
# create
|
|
p_cr = sub.add_parser("create", help="Crea un notebook vacio via API REST")
|
|
p_cr.add_argument("notebook", help="Ruta del notebook a crear")
|
|
p_cr.add_argument("--kernel", default="python3", help="Nombre del kernel (default: python3)")
|
|
p_cr.add_argument("--force", action="store_true", help="Sobreescribir si ya existe")
|
|
add_common(p_cr)
|
|
|
|
# batch
|
|
p_bt = sub.add_parser("batch", help="Anade multiples celdas en una sola conexion WebSocket")
|
|
p_bt.add_argument("notebook", help="Ruta del notebook")
|
|
p_bt.add_argument(
|
|
"--from",
|
|
dest="cells_source",
|
|
default="-",
|
|
metavar="FILE",
|
|
help="Archivo JSON con las celdas, o '-' para leer de stdin (default: -)",
|
|
)
|
|
add_common(p_bt)
|
|
|
|
return parser
|
|
|
|
|
|
def main() -> None:
|
|
parser = _build_parser()
|
|
args = parser.parse_args()
|
|
|
|
if args.command == "append-code":
|
|
result = jupyter_append_code(args.notebook, args.source, args.server, args.token)
|
|
elif args.command == "append-markdown":
|
|
result = jupyter_append_markdown(args.notebook, args.source, args.server, args.token)
|
|
elif args.command == "insert":
|
|
result = jupyter_insert_cell(
|
|
args.notebook, args.index, args.source, args.cell_type, args.server, args.token
|
|
)
|
|
elif args.command == "edit":
|
|
result = jupyter_edit_cell(args.notebook, args.index, args.source, args.server, args.token)
|
|
elif args.command == "delete":
|
|
result = jupyter_delete_cell(args.notebook, args.index, args.server, args.token)
|
|
elif args.command == "create":
|
|
result = jupyter_create_notebook(
|
|
args.notebook,
|
|
kernel_name=args.kernel,
|
|
server_url=args.server,
|
|
token=args.token,
|
|
force=args.force,
|
|
)
|
|
elif args.command == "batch":
|
|
if args.cells_source == "-":
|
|
raw = sys.stdin.read()
|
|
else:
|
|
with open(args.cells_source, "r", encoding="utf-8") as f:
|
|
raw = f.read()
|
|
cells = json.loads(raw)
|
|
result = jupyter_batch_write(args.notebook, cells, args.server, args.token)
|
|
else:
|
|
parser.print_help()
|
|
return
|
|
|
|
print(json.dumps(result, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|