first steps to a post-transform step for the AST

Author: choldgraf

jupyter_sphinx/__init__.py

@@ -6,16 +6,21 @@
 import ipywidgets
 import os
 from sphinx.util.fileutil import copy_asset
+from sphinx.errors import ExtensionError
 from IPython.lib.lexers import IPythonTracebackLexer, IPython3Lexer
 
 from .ast import (
     JupyterCell,
     JupyterCellNode,
+    CellInputNode,
+    CellOutputNode,
+    CellOutputBundleNode,
     JupyterKernelNode,
     JupyterWidgetViewNode,
     JupyterWidgetStateNode,
     WIDGET_VIEW_MIMETYPE,
     jupyter_download_role,
+    CellOutputsToNodes,
 )
 from .execute import JupyterKernel, ExecuteJupyterCells
 from .thebelab import ThebeButton, ThebeButtonNode, ThebeOutputNode, ThebeSourceNode
@@ -34,6 +39,11 @@
 def skip(self, node):
     raise docutils.nodes.SkipNode
 
+# Used for nodes that should be gone by rendering time (OutputMimeBundleNode)
+def halt(self, node):
+    raise ExtensionError((f"Rendering encountered a node type that should "
+                         "have been removed before rendering: {type(node)}"))
+
 # Renders the children of a container
 render_container = (
     lambda self, node: self.visit_container(node),
@@ -45,22 +55,26 @@ def visit_container_html(self, node):
     self.body.append(node.visit_html())
     self.visit_container(node)
 
+
 def depart_container_html(self, node):
     self.depart_container(node)
     self.body.append(node.depart_html())
 
+
 # Used to render an element node as HTML
 def visit_element_html(self, node):
     self.body.append(node.html())
     raise docutils.nodes.SkipNode
 
+
 # Used to render the ThebeSourceNode conditionally for non-HTML builders
 def visit_thebe_source(self, node):
     if node["hide_code"]:
         raise docutils.nodes.SkipNode
     else:
         self.visit_container(node)
 
+
 render_thebe_source = (
     visit_thebe_source,
     lambda self, node: self.depart_container(node),
@@ -170,15 +184,28 @@ def setup(app):
         man=(skip, None),
     )
 
-    # JupyterCellNode is a container that holds the input and
-    # any output, so we render it as a container.
+    # Register our container nodes, these should behave just like a regular container
+    for node in [JupyterCellNode, CellInputNode, CellOutputNode]:
+        app.add_node(
+            node,
+            override=True,
+            html=(render_container),
+            latex=(render_container),
+            textinfo=(render_container),
+            text=(render_container),
+            man=(render_container),
+        )
+
+    # Register the output bundle node.
+    # No translators should touch this node because we'll replace it in a post-transform
     app.add_node(
-        JupyterCellNode,
-        html=render_container,
-        latex=render_container,
-        textinfo=render_container,
-        text=render_container,
-        man=render_container,
+        CellOutputBundleNode,
+        override=True,
+        html=(halt, None),
+        latex=(halt, None),
+        textinfo=(halt, None),
+        text=(halt, None),
+        man=(halt, None),
     )
 
     # JupyterWidgetViewNode holds widget view JSON,
@@ -242,6 +269,7 @@ def setup(app):
     app.add_role("jupyter-download:notebook", jupyter_download_role)
     app.add_role("jupyter-download:script", jupyter_download_role)
     app.add_transform(ExecuteJupyterCells)
+    app.add_post_transform(CellOutputsToNodes)
 
     # For syntax highlighting
     app.add_lexer("ipythontb", IPythonTracebackLexer())

jupyter_sphinx/ast.py

@@ -5,9 +5,11 @@
 
 import docutils
 from docutils.parsers.rst import Directive, directives
-from docutils.nodes import math_block
+from docutils.nodes import math_block, image
 from sphinx.util import parselinenos
 from sphinx.addnodes import download_reference
+from sphinx.transforms import SphinxTransform
+from sphinx.environment.collectors.asset import ImageCollector
 
 import ipywidgets.embed
 import nbconvert
@@ -127,20 +129,24 @@ def run(self):
         else:
             hl_lines = []
 
-        return [
-            JupyterCellNode(
-                "",
-                docutils.nodes.literal_block(text="\n".join(content)),
-                hide_code=("hide-code" in self.options),
-                hide_output=("hide-output" in self.options),
-                code_below=("code-below" in self.options),
-                linenos=("linenos" in self.options),
-                linenostart=(self.options.get("lineno-start")),
-                emphasize_lines=hl_lines,
-                raises=self.options.get("raises"),
-                stderr=("stderr" in self.options),
-            )
-        ]
+        # A top-level placeholder for our cell
+        cell_node = JupyterCellNode(
+            hide_code=("hide-code" in self.options),
+            hide_output=("hide-output" in self.options),
+            code_below=("code-below" in self.options),
+            linenos=("linenos" in self.options),
+            linenostart=(self.options.get("lineno-start")),
+            emphasize_lines=hl_lines,
+            raises=self.options.get("raises"),
+            stderr=("stderr" in self.options),
+            classes=["jupyter_cell"],
+        )
+
+        # Add the input section of the cell, we'll add output at execution time
+        cell_input = CellInputNode(classes=["cell_input"])
+        cell_input += docutils.nodes.literal_block(text="\n".join(content))
+        cell_node += cell_input
+        return [cell_node]
 
 
 class JupyterCellNode(docutils.nodes.container):
@@ -151,6 +157,28 @@ class JupyterCellNode(docutils.nodes.container):
     """
 
 
+class CellInputNode(docutils.nodes.container):
+    """Represent an input cell in the Sphinx AST."""
+
+    def __init__(self, rawsource="", *children, **attributes):
+        super().__init__("", **attributes)
+
+
+class CellOutputNode(docutils.nodes.container):
+    """Represent an output cell in the Sphinx AST."""
+
+    def __init__(self, rawsource="", *children, **attributes):
+        super().__init__("", **attributes)
+
+
+class CellOutputBundleNode(docutils.nodes.container):
+    """Represent a MimeBundle in the Sphinx AST, to be transformed later."""
+
+    def __init__(self, outputs, rawsource="", *children, **attributes):
+        self.outputs = outputs
+        super().__init__("", **attributes)
+
+
 class JupyterKernelNode(docutils.nodes.Element):
     """Inserted into doctree whenever a JupyterKernel directive is encountered.
 
@@ -199,12 +227,12 @@ def html(self):
         )
 
 
-def cell_output_to_nodes(cell, data_priority, write_stderr, dir, thebe_config):
+def cell_output_to_nodes(outputs, data_priority, write_stderr, dir, thebe_config):
     """Convert a jupyter cell with outputs and filenames to doctree nodes.
 
     Parameters
     ----------
-    cell : jupyter cell
+    outputs : a list of outputs from a Jupyter cell
     data_priority : list of mime types
         Which media types to prioritize.
     write_stderr : bool
@@ -214,9 +242,14 @@ def cell_output_to_nodes(cell, data_priority, write_stderr, dir, thebe_config):
         to the source folder prefixed with ``/``.
     thebe_config: dict
         Thebelab configuration object or None
+
+    Returns
+    -------
+    to_add : list of docutils nodes
+        Each output, converted into a docutils node.
     """
     to_add = []
-    for _, output in enumerate(cell.get("outputs", [])):
+    for output in outputs:
         output_type = output["output_type"]
         if output_type == "stream":
             if output["name"] == "stderr":
@@ -270,6 +303,9 @@ def cell_output_to_nodes(cell, data_priority, write_stderr, dir, thebe_config):
                 continue
             data = output["data"][mime_type]
             if mime_type.startswith("image"):
+                ####################################
+                # TODO: Figure out how to handle either inline or absolute image paths
+
                 # Sphinx treats absolute paths as being rooted at the source
                 # directory, so make a relative path, which Sphinx treats
                 # as being relative to the current working directory.
@@ -326,32 +362,37 @@ def cell_output_to_nodes(cell, data_priority, write_stderr, dir, thebe_config):
 def attach_outputs(output_nodes, node, thebe_config, cm_language):
     if not node.attributes["hide_code"]:  # only add css if code is displayed
         node.attributes["classes"] = ["jupyter_container"]
+
+    input_node = _return_first_node_type(node, CellInputNode)
+    outputbundle_node = _return_first_node_type(node, CellOutputBundleNode)
+    output_node = CellOutputNode(classes=["cell_output"])
     if thebe_config:
-        source = node.children[0]
+        # Move the source from the input node into the thebe_source node
+        source = input_node.children.pop(0)
         thebe_source = ThebeSourceNode(
             hide_code=node.attributes["hide_code"],
             code_below=node.attributes["code_below"],
             language=cm_language,
         )
         thebe_source.children = [source]
-
-        node.children = [thebe_source]
+        input_node.children = [thebe_source]
 
         if not node.attributes["hide_output"]:
             thebe_output = ThebeOutputNode()
             thebe_output.children = output_nodes
-            if node.attributes["code_below"]:
-                node.children = [thebe_output] + node.children
-            else:
-                node.children = node.children + [thebe_output]
+            output_node += thebe_output
     else:
         if node.attributes["hide_code"]:
-            node.children = []
+            node.children.pop(0)
         if not node.attributes["hide_output"]:
-            if node.attributes["code_below"]:
-                node.children = output_nodes + node.children
-            else:
-                node.children = node.children + output_nodes
+            output_node.children = output_nodes
+
+    # Now replace the bundle with our OutputNode
+    outputbundle_node.replace_self(output_node)
+
+    # Swap inputs and outputs if we want the code below
+    if node.attributes["code_below"]:
+        node.children = node.children[::-1]
 
 
 def jupyter_download_role(name, rawtext, text, lineno, inliner):
@@ -373,3 +414,50 @@ def get_widgets(notebook):
         # Don't catch KeyError, as it's a bug if 'widgets' does
         # not contain 'WIDGET_STATE_MIMETYPE'
         return None
+
+
+class CellOutputsToNodes(SphinxTransform):
+    """Use the builder context to transform a CellOutputNode into Sphinx nodes."""
+
+    default_priority = 700
+
+    def apply(self):
+        thebe_config = self.config.jupyter_sphinx_thebelab_config
+
+        for cell_node in self.document.traverse(JupyterCellNode):
+            output_bundle_node = _return_first_node_type(
+                cell_node, CellOutputBundleNode
+            )
+            # Create doctree nodes for cell outputs.
+            output_nodes = cell_output_to_nodes(
+                output_bundle_node.outputs,
+                self.config.jupyter_execute_data_priority,
+                bool(cell_node.attributes["stderr"]),
+                sphinx_abs_dir(self.env),
+                thebe_config,
+            )
+            # Remove the outputbundlenode and we'll attach the outputs next
+            attach_outputs(output_nodes, cell_node, thebe_config, cell_node.cm_language)
+
+        # Image collect extra nodes from cell outputs that we need to process
+        for node in self.document.traverse(image):
+            # If the image node has `candidates` then it's already been processed
+            # as in-line markdown, so skip it
+            if "candidates" in node:
+                continue
+            col = ImageCollector()
+            col.process_doc(self.app, node)
+
+
+def _return_first_node_type(node, node_type):
+    found_nodes = list(node.traverse(node_type))
+    if len(found_nodes) == 0:
+        raise ValueError(f"Found no nodes of type {node_type} in node {node}")
+    if len(found_nodes) > 1:
+        raise ValueError(
+            (
+                f"Found more than one nodes of type {node_type} in node {node}. "
+                "only return the first instance"
+            )
+        )
+    return found_nodes[0]

jupyter_sphinx/css/jupyter-sphinx.css

@@ -58,7 +58,7 @@ div.jupyter_container {
 div.jupyter_container div.highlight {
   background-color: #f7f7f7; /* for haiku */
 }
-div.jupyter_container {
+div.jupyter_container * {
     padding: 0;
     margin: 0;
 }

jupyter_sphinx/execute.py

@@ -27,6 +27,7 @@
 )
 from .ast import (
     JupyterCellNode,
+    CellOutputBundleNode,
     JupyterKernelNode,
     cell_output_to_nodes,
     JupyterWidgetStateNode,
@@ -200,7 +201,7 @@ def apply(self):
             # Add code cell CSS class
             for node in nodes:
                 source = node.children[0]
-                source.attributes["classes"] = ["code_cell"]
+                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
@@ -211,17 +212,12 @@ def apply(self):
                 cm_language = notebook.metadata.language_info.codemirror_mode.name
             except AttributeError:
                 cm_language = notebook.metadata.kernelspec.language
+            for node in nodes:
+                node.cm_language = cm_language
 
             # Add doctree nodes for cell outputs.
             for node, cell in zip(nodes, notebook.cells):
-                output_nodes = cell_output_to_nodes(
-                    cell,
-                    self.config.jupyter_execute_data_priority,
-                    bool(node.attributes["stderr"]),
-                    sphinx_abs_dir(self.env),
-                    thebe_config,
-                )
-                attach_outputs(output_nodes, node, thebe_config, cm_language)
+                node += CellOutputBundleNode(cell.outputs)
 
             if contains_widgets(notebook):
                 doctree.append(JupyterWidgetStateNode(state=get_widgets(notebook)))