"""Execution and managing kernels.""" import os import warnings from logging import Logger from pathlib import Path import nbconvert from docutils.parsers.rst import Directive, directives from nbconvert.preprocessors import ExtractOutputPreprocessor from nbconvert.writers import FilesWriter from sphinx.errors import ExtensionError from sphinx.transforms import SphinxTransform if nbconvert.version_info < (6,): from nbconvert.preprocessors.execute import executenb else: from nbclient.client import execute as executenb import traitlets # Workaround of https://github.com/ipython/traitlets/issues/606 if traitlets.version_info < (5, 1): class LoggerAdapterWrapper(Logger): """Wrap a logger adapter, while pretending to be a logger.""" def __init__(self, wrapped): self._wrapped = wrapped def __getattribute__(self, attr): if attr == "_wrapped": return object.__getattribute__(self, attr) return self._wrapped.__getattribute__(attr) else: def LoggerAdapterWrapper(logger_adapter): return logger_adapter import nbformat import jupyter_sphinx as js from .ast import ( CellOutputNode, JupyterCellNode, JupyterKernelNode, JupyterWidgetStateNode, apply_styling, cell_output_to_nodes, get_widgets, ) from .thebelab import ThebeButtonNode, add_thebelab_library from .utils import ( blank_nb, default_notebook_names, output_directory, sphinx_abs_dir, split_on, ) class JupyterKernel(Directive): """Specify a new Jupyter Kernel. Arguments --------- kernel_name : str (optional) The name of the kernel in which to execute future Jupyter cells, as reported by executing 'jupyter kernelspec list' on the command line. Options ------- id : str An identifier for *this kernel instance*. Used to name any output files generated when executing the Jupyter cells (e.g. images produced by cells, or a script containing the cell inputs). Content ------- None """ optional_arguments = 1 final_argument_whitespace = False has_content = False option_spec = {"id": directives.unchanged} def run(self): return [ JupyterKernelNode( "", kernel_name=self.arguments[0].strip() if self.arguments else "", kernel_id=self.options.get("id", "").strip(), ) ] # Doctree transformations class ExecuteJupyterCells(SphinxTransform): """Execute code cells in Jupyter kernels. Traverses the doctree to find JupyterKernel and JupyterCell nodes, then executes the code in the JupyterCell nodes in sequence, starting a new kernel every time a JupyterKernel node is encountered. The output from each code cell is inserted into the doctree. """ # Beginning of main transforms. Not 100% sure it's the correct time. default_priority = 400 def apply(self): doctree = self.document docname_path = Path(self.env.docname) doc_dir_relpath = docname_path.parent # relative to src dir docname = docname_path.name default_kernel = self.config.jupyter_execute_default_kernel default_names = default_notebook_names(docname) thebe_config = self.config.jupyter_sphinx_thebelab_config linenos_config = self.config.jupyter_sphinx_linenos continue_linenos = self.config.jupyter_sphinx_continue_linenos # Check if we have anything to execute. if not next(doctree.findall(JupyterCellNode), False): return if thebe_config: # Add the button at the bottom if it is not present if not next(doctree.findall(ThebeButtonNode), False): doctree.append(ThebeButtonNode()) add_thebelab_library(doctree, self.env) js.logger.info(f"executing {docname}") output_dir = Path(output_directory(self.env)) / doc_dir_relpath # Start new notebook whenever a JupyterKernelNode is encountered jupyter_nodes = (JupyterCellNode, JupyterKernelNode) nodes_by_notebook = split_on( lambda n: isinstance(n, JupyterKernelNode), list(doctree.findall(lambda n: isinstance(n, jupyter_nodes))), ) for first, *nodes in nodes_by_notebook: if isinstance(first, JupyterKernelNode): kernel_name = first["kernel_name"] or default_kernel file_name = first["kernel_id"] or next(default_names) else: nodes = (first, *nodes) kernel_name = default_kernel file_name = next(default_names) # Add empty placeholder cells for non-executed nodes so nodes # and cells can be zipped and the provided input/output # can be inserted later notebook = execute_cells( kernel_name, [ nbformat.v4.new_code_cell(node.astext() if node["execute"] else "") for node in nodes ], self.config.jupyter_execute_kwargs, ) # Raise error if cells raised exceptions and were not marked as doing so for node, cell in zip(nodes, notebook.cells): errors = [ output for output in cell.outputs if output["output_type"] == "error" ] allowed_errors = node.attributes.get("raises") or [] raises_provided = node.attributes["raises"] is not None if ( raises_provided and not allowed_errors ): # empty 'raises': suppress all errors pass elif errors and not any(e["ename"] in allowed_errors for e in errors): raise ExtensionError( "Cell raised uncaught exception:\n{}".format( "\n".join(errors[0]["traceback"]) ) ) # Raise error if cells print to stderr for node, cell in zip(nodes, notebook.cells): stderr = [ output for output in cell.outputs if output["output_type"] == "stream" and output["name"] == "stderr" ] if stderr and not node.attributes["stderr"]: js.logger.warning( "Cell printed to stderr:\n{}".format(stderr[0]["text"]) ) # Insert input/output into placeholders for non-executed cells for node, cell in zip(nodes, notebook.cells): if not node["execute"]: cell.source = node.children[0].astext() if len(node.children) == 2: output = nbformat.v4.new_output("stream") output.text = node.children[1].astext() cell.outputs = [output] node.children.pop() try: lexer = notebook.metadata.language_info.pygments_lexer except AttributeError: lexer = notebook.metadata.kernelspec.language # Highlight the code cells now that we know what language they are for node in nodes: source = node.children[0].children[0] source.attributes["language"] = lexer # Add line numbering linenostart = 1 for node in nodes: # The literal_block node with the source source = node.children[0].children[0] nlines = source.rawsource.count("\n") + 1 show_numbering = ( linenos_config or source["linenos"] or source["linenostart"] ) if show_numbering: source["linenos"] = True if source["linenostart"]: linenostart = source["linenostart"] if source["linenostart"] or continue_linenos: source["highlight_args"] = {"linenostart": linenostart} else: linenostart = 1 linenostart += nlines hl_lines = node["emphasize_lines"] if hl_lines: highlight_args = source.setdefault("highlight_args", {}) highlight_args["hl_lines"] = hl_lines # Add code cell CSS class for node in nodes: source = node.children[0] source.attributes["classes"].append("code_cell") # Write certain cell outputs (e.g. images) to separate files, and # modify the metadata of the associated cells in 'notebook' to # include the path to the output file. write_notebook_output( notebook, str(output_dir), file_name, self.env.docname ) try: cm_language = notebook.metadata.language_info.codemirror_mode.name except AttributeError: cm_language = notebook.metadata.kernelspec.language for node in nodes: node.attributes["cm_language"] = cm_language # Add doctree nodes for cell outputs. for node, cell in zip(nodes, notebook.cells): # Add the outputs as children output = CellOutputNode(classes=["cell_output"]) output.children = cell_output_to_nodes( cell.outputs, bool(node.attributes["stderr"]), sphinx_abs_dir(self.env), thebe_config, ) node += output apply_styling(node, thebe_config) if contains_widgets(notebook): doctree.append(JupyterWidgetStateNode(state=get_widgets(notebook))) # Roles def execute_cells(kernel_name, cells, execute_kwargs): """Execute Jupyter cells in the specified kernel and return the notebook.""" notebook = blank_nb(kernel_name) notebook.cells = cells # Modifies 'notebook' in-place try: executenb(notebook, **execute_kwargs) except Exception as e: raise ExtensionError("Notebook execution failed", orig_exc=e) return notebook def write_notebook_output(notebook, output_dir, notebook_name, location=None): """Extract output from notebook cells and write to files in output_dir. This also modifies 'notebook' in-place, adding metadata to each cell that maps output mime-types to the filenames the output was saved under. """ resources = dict(unique_key=os.path.join(output_dir, notebook_name), outputs={}) # Modifies 'resources' in-place ExtractOutputPreprocessor().preprocess(notebook, resources) # Write the cell outputs to files where we can (images and PDFs), # as well as the notebook file. FilesWriter(build_directory=output_dir).write( nbformat.writes(notebook), resources, os.path.join(output_dir, notebook_name + ".ipynb"), ) exporter = nbconvert.exporters.ScriptExporter(log=LoggerAdapterWrapper(js.logger)) with warnings.catch_warnings(): # See https://github.com/jupyter/nbconvert/issues/1388 warnings.simplefilter("ignore", DeprecationWarning) contents, resources = exporter.from_notebook_node(notebook) notebook_file = notebook_name + resources["output_extension"] output_dir = Path(output_dir) # utf-8 is the de-facto standard encoding for notebooks. (output_dir / notebook_file).write_text(contents, encoding="utf8") def contains_widgets(notebook): widgets = get_widgets(notebook) return widgets and widgets["state"]