68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
import argparse
|
|
import sys
|
|
from jupyter_client import KernelManager
|
|
import nbformat
|
|
from nbconvert import HTMLExporter
|
|
from nbclient import NotebookClient
|
|
from nbclient.exceptions import CellExecutionError
|
|
import json
|
|
import logging
|
|
|
|
|
|
# Parse input arguments.
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('document_id')
|
|
args = parser.parse_args()
|
|
|
|
# Set up a logger that writes to stderr.
|
|
logging.basicConfig(level='INFO')
|
|
logger = logging.getLogger('obsidian-jupyter')
|
|
logger.info('started server for document %s', args.document_id)
|
|
|
|
# Create a notebook and kernel.
|
|
cell = nbformat.v4.new_code_cell()
|
|
nb = nbformat.v4.new_notebook(cells=[cell])
|
|
km = KernelManager()
|
|
client = NotebookClient(nb, km)
|
|
|
|
try:
|
|
# Respond to each request.
|
|
for request in sys.stdin:
|
|
# Load the request and generate a response with matching id.
|
|
logger.info('received request: %s', request)
|
|
request = json.loads(request)
|
|
request_body = request['body']
|
|
response = {
|
|
'id': request['id'],
|
|
}
|
|
# Execute a cell.
|
|
if request_body['command'] == 'execute':
|
|
cell['source'] = request_body['source']
|
|
try:
|
|
nb = client.execute(nb)
|
|
except CellExecutionError as ex:
|
|
logger.info('cell failed to execute: %s', ex)
|
|
html_exporter = HTMLExporter(template_name='basic')
|
|
(response_body, resources) = html_exporter.from_notebook_node(nb)
|
|
elif request_body['command'] == 'restart_kernel':
|
|
km.restart_kernel()
|
|
response_body = ''
|
|
else:
|
|
logger.error('unrecognised command: %s', request_body['command'])
|
|
response_body = ''
|
|
|
|
# Pass the response back.
|
|
response['body'] = response_body
|
|
response = json.dumps(response)
|
|
sys.stdout.write(response + '\n')
|
|
logger.info('sent response: %s', response)
|
|
sys.stdout.flush()
|
|
sys.stderr.flush()
|
|
finally:
|
|
# Clean up the kernel.
|
|
if km.is_alive:
|
|
logger.info('shutting down kernel...')
|
|
km.shutdown_kernel()
|
|
|
|
logger.info('exiting...')
|