Execute preprocess cells

nbconvert/preprocessors/execute.py

@@ -3,7 +3,10 @@
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
+from typing import Optional
+from nbformat import NotebookNode
 from nbclient import NotebookClient, execute as _execute
+from nbclient.util import run_sync
 # Backwards compatability for imported name
 from nbclient.exceptions import CellExecutionError
 
@@ -31,12 +34,27 @@ def __init__(self, **kw):
         Preprocessor.__init__(self, nb=nb, **kw)
         NotebookClient.__init__(self, nb, **kw)
 
+    def _check_assign_resources(self, resources):
+        if resources or not hasattr(self, 'resources'):
+            self.resources = resources
+
     def preprocess(self, nb, resources=None, km=None):
         """
         Preprocess notebook executing each code cell.
 
         The input argument *nb* is modified in-place.
 
+        Note that this function recalls NotebookClient.__init__, which may look wrong.
+        However since the preprocess call acts line an init on exeuction state it's expected.
+        Therefore, we need to capture it here again to properly reset because traitlet 
+        assignments are not passed. There is a risk if traitlets apply any side effects for
+        dual init.
+        The risk should be manageable, and this approach minimizes side-effects relative
+        to other alternatives.
+
+        One alternative but rejected implementation would be to copy the client's init internals
+        which has already gotten out of sync with nbclient 0.5 release before nbcovnert 6.0 released.
+
         Parameters
         ----------
         nb : NotebookNode
@@ -56,11 +74,73 @@ def preprocess(self, nb, resources=None, km=None):
         resources : dictionary
             Additional resources used in the conversion process.
         """
-        # Copied from NotebookClient init :/
-        self.nb = nb
-        self.km = km
-        if resources:
-            self.resources = resources
-        self.reset_execution_trackers()
+        NotebookClient.__init__(self, nb, km)
+        self._check_assign_resources(resources)
         self.execute()
-        return nb, resources
+        return self.nb, self.resources
+
+    async def async_execute_cell(
+            self,
+            cell: NotebookNode,
+            cell_index: int,
+            execution_count: Optional[int] = None,
+            store_history: bool = False) -> NotebookNode:
+        """
+        Executes a single code cell.
+
+        Overwrites NotebookClient's version of this method to allow for preprocess_cell calls.
+
+        Parameters
+        ----------
+        cell : nbformat.NotebookNode
+            The cell which is currently being processed.
+        cell_index : int
+            The position of the cell within the notebook object.
+        execution_count : int
+            The execution count to be assigned to the cell (default: Use kernel response)
+        store_history : bool
+            Determines if history should be stored in the kernel (default: False).
+            Specific to ipython kernels, which can store command histories.
+
+        Returns
+        -------
+        output : dict
+            The execution output payload (or None for no output).
+
+        Raises
+        ------
+        CellExecutionError
+            If execution failed and should raise an exception, this will be raised
+            with defaults about the failure.
+
+        Returns
+        -------
+        cell : NotebookNode
+            The cell which was just processed.
+        """
+        # Copied and intercepted to allow for custom preprocess_cell contracts to be fullfilled
+        self.store_history = store_history
+        cell, resources = self.preprocess_cell(cell, self.resources, cell_index)
+        if execution_count:
+            cell['execution_count'] = execution_count
+        return cell, resources
+
+    def preprocess_cell(self, cell, resources, index, **kwargs):
+        """
+        Override if you want to apply some preprocessing to each cell.
+        Must return modified cell and resource dictionary.
+
+        Parameters
+        ----------
+        cell : NotebookNode cell
+            Notebook cell being processed
+        resources : dictionary
+            Additional resources used in the conversion process.  Allows
+            preprocessors to pass variables into the Jinja engine.
+        index : int
+            Index of the cell being processed
+        """
+        self._check_assign_resources(resources)
+        # Because nbclient is an async library, we need to wrap the parent async call to generate a syncronous version.
+        cell = run_sync(NotebookClient.async_execute_cell)(self, cell, index, store_history=self.store_history)
+        return cell, self.resources

nbconvert/preprocessors/tests/test_execute.py

@@ -11,6 +11,8 @@
 import pytest
 import nbformat
 
+from copy import deepcopy
+
 from ..execute import ExecutePreprocessor, executenb
 
 
@@ -58,19 +60,41 @@ def test_basic_execution():
     fname = os.path.join(os.path.dirname(__file__), 'files', 'HelloWorld.ipynb')
     with open(fname) as f:
         input_nb = nbformat.read(f, 4)
-        output_nb, _ = preprocessor.preprocess(input_nb)
+        output_nb, _ = preprocessor.preprocess(deepcopy(input_nb))
     assert_notebooks_equal(input_nb, output_nb)
 
+
 def test_executenb():
     fname = os.path.join(os.path.dirname(__file__), 'files', 'HelloWorld.ipynb')
     with open(fname) as f:
         input_nb = nbformat.read(f, 4)
         with pytest.warns(FutureWarning):
-            output_nb = executenb(input_nb)
+            output_nb = executenb(deepcopy(input_nb))
     assert_notebooks_equal(input_nb, output_nb)
 
+
 def test_populate_language_info():
     preprocessor = ExecutePreprocessor(kernel_name="python")
     nb = nbformat.v4.new_notebook()  # Certainly has no language_info.
-    nb, _ = preprocessor.preprocess(nb, resources={})
+    preprocessor.preprocess(nb, resources={})
+    # Should mutate input
     assert 'language_info' in nb.metadata  # See that a basic attribute is filled in
+
+
+def test_preprocess_cell():
+    class CellReplacer(ExecutePreprocessor):
+        def preprocess_cell(self, cell, resources, index, **kwargs):
+            cell.source = "print('Ignored')"
+            return super().preprocess_cell(cell, resources, index, **kwargs)
+
+    preprocessor = CellReplacer()
+    fname = os.path.join(os.path.dirname(__file__), 'files', 'HelloWorld.ipynb')
+    with open(fname) as f:
+        input_nb = nbformat.read(f, 4)
+        output_nb, _ = preprocessor.preprocess(deepcopy(input_nb))
+    expected_nb = deepcopy(input_nb)
+    for cell in expected_nb.cells:
+        cell.source = "print('Ignored')"
+        for output in cell.outputs:
+            output.text = 'Ignored\n'
+    assert_notebooks_equal(expected_nb, output_nb)