Merge nodes from jupyter-input and jupyter-output

Author: willsimmons1465

jupyter_sphinx/__init__.py

@@ -23,6 +23,7 @@
     JupyterWidgetStateNode,
     WIDGET_VIEW_MIMETYPE,
     JupyterDownloadRole,
+    CombineCellInputOutput,
     CellOutputsToNodes,
 )
 from .execute import JupyterKernel, ExecuteJupyterCells
@@ -275,6 +276,7 @@ def setup(app):
     app.add_role("jupyter-download:notebook", JupyterDownloadRole())
     app.add_role("jupyter-download:nb", JupyterDownloadRole())
     app.add_role("jupyter-download:script", JupyterDownloadRole())
+    app.add_transform(CombineCellInputOutput)
     app.add_transform(ExecuteJupyterCells)
     app.add_transform(CellOutputsToNodes)
 

jupyter_sphinx/ast.py

@@ -12,6 +12,7 @@
 from sphinx.addnodes import download_reference
 from sphinx.transforms import SphinxTransform
 from sphinx.environment.collectors.asset import ImageCollector
+from sphinx.errors import ExtensionError
 
 import ipywidgets.embed
 import nbconvert
@@ -271,7 +272,6 @@ def run(self):
             emphasize_lines=[],
             raises=False,
             stderr=False,
-            classes=["jupyter_cell"],
         )
 
         # Add a blank input and the given output to the cell
@@ -574,6 +574,39 @@ def get_widgets(notebook):
         return None
 
 
+class CombineCellInputOutput(SphinxTransform):
+    """Merge nodes from CellOutput with the preceding CellInput node."""
+
+    default_priority = 120
+
+    def apply(self):
+        moved_outputs = set()
+
+        for cell_node in self.document.traverse(JupyterCellNode):
+            if cell_node.attributes["execute"] == False:
+                if cell_node.attributes["hide_code"] == False:
+                    # Cell came from jupyter-input
+                    sibling = cell_node.next_node(descend=False, siblings=True)
+                    if (
+                        isinstance(sibling, JupyterCellNode)
+                        and sibling.attributes["execute"] == False
+                        and sibling.attributes["hide_code"] == True
+                    ):
+                        # Sibling came from jupyter-output, so we merge
+                        cell_node += sibling.children[1]
+                        cell_node.attributes["hide_output"] = False
+                        moved_outputs.update({sibling})
+                else:
+                    # Call came from jupyter-output
+                    if cell_node not in moved_outputs:
+                        raise ExtensionError(
+                            "Found a jupyter-output node without a preceding jupyter-input"
+                        )
+
+        for output_node in moved_outputs:
+            output_node.replace_self([])
+
+
 class CellOutputsToNodes(SphinxTransform):
     """Use the builder context to transform a CellOutputNode into Sphinx nodes."""
 

tests/test_execute.py

@@ -682,15 +682,29 @@ def test_input_cell_linenos(doctree):
 
 def test_output_cell(doctree):
     source = """
+    .. jupyter-input::
+
+        3 + 2
+
     .. jupyter-output::
 
         4
     """
     tree = doctree(source)
     (cell,) = tree.traverse(JupyterCellNode)
-    (celloutput,) = cell.children
+    (cellinput, celloutput,) = cell.children
+    assert cellinput.children[0].rawsource.strip() == "3 + 2"
     assert celloutput.children[0].rawsource.strip() == "4"
 
+def test_output_only_error(doctree):
+    source = """
+    .. jupyter-output::
+
+        4
+    """
+    with pytest.raises(ExtensionError):
+        tree = doctree(source)
+
 def test_multiple_directives(doctree):
     source = """
     .. jupyter-execute::
@@ -706,11 +720,10 @@ def test_multiple_directives(doctree):
         5
     """
     tree = doctree(source)
-    (ex, jin, jout) = tree.traverse(JupyterCellNode)
+    (ex, jin) = tree.traverse(JupyterCellNode)
     (ex_in, ex_out) = ex.children
-    (jin_in, _) = jin.children
-    (jout_out,) = jout.children
+    (jin_in, jin_out) = jin.children
     assert ex_in.children[0].rawsource.strip() == "2 + 2"
     assert ex_out.children[0].rawsource.strip() == "4"
     assert jin_in.children[0].rawsource.strip() == "3 + 3"
-    assert jout_out.children[0].rawsource.strip() == "5"
+    assert jin_out.children[0].rawsource.strip() == "5"