From 114d66e43f16c7c2fc41ccfacf96c0deedadf0e3 Mon Sep 17 00:00:00 2001
From: Matthew Seal <mseal007@gmail.com>
Date: Sun, 28 Oct 2018 16:56:01 -0700
Subject: [PATCH] Refactored execute preprocessor to have a process_message
 function for each kernel message

---
 nbconvert/preprocessors/execute.py | 94 ++++++++++++++++--------------
 1 file changed, 49 insertions(+), 45 deletions(-)

diff --git a/nbconvert/preprocessors/execute.py b/nbconvert/preprocessors/execute.py
index fe80c5d66..31644c5bc 100644
--- a/nbconvert/preprocessors/execute.py
+++ b/nbconvert/preprocessors/execute.py
@@ -446,7 +446,7 @@ def run_cell(self, cell, cell_index=0):
         self.log.debug("Executing cell:\n%s", cell.source)
         exec_reply = self._wait_for_reply(msg_id, cell)
 
-        outs = cell.outputs = []
+        cell.outputs = []
 
         while True:
             try:
@@ -466,55 +466,59 @@ def run_cell(self, cell, cell_index=0):
                 # not an output from our execution
                 continue
 
-            msg_type = msg['msg_type']
-            self.log.debug("output: %s", msg_type)
-            content = msg['content']
-
-            # set the prompt number for the input and the output
-            if 'execution_count' in content:
-                cell['execution_count'] = content['execution_count']
-
-            if msg_type == 'status':
-                if content['execution_state'] == 'idle':
-                    break
-                else:
-                    continue
-            elif msg_type == 'execute_input':
-                continue
-            elif msg_type == 'clear_output':
-                outs[:] = []
-                # clear display_id mapping for this cell
-                for display_id, cell_map in self._display_id_map.items():
-                    if cell_index in cell_map:
-                        cell_map[cell_index] = []
-                continue
-            elif msg_type.startswith('comm'):
-                continue
-
-            display_id = None
-            if msg_type in {'execute_result', 'display_data', 'update_display_data'}:
-                display_id = msg['content'].get('transient', {}).get('display_id', None)
-                if display_id:
-                    self._update_display_id(display_id, msg)
-                if msg_type == 'update_display_data':
-                    # update_display_data doesn't get recorded
-                    continue
-
+            if not self.process_message(msg, cell, cell_index):
+                break
+
+        return exec_reply, cell.outputs
+
+    def process_message(self, msg, cell, cell_index):
+        '''
+        Returns None if execution should be halted.
+        '''
+        msg_type = msg['msg_type']
+        self.log.debug("msg_type: %s", msg_type)
+        content = msg['content']
+        self.log.debug("content: %s", content)
+
+        # Default to our input as the "result" of processing the message
+        result = msg
+
+        display_id = content.get('transient', {}).get('display_id', None)
+        if display_id and msg_type in {'execute_result', 'display_data', 'update_display_data'}:
+            self._update_display_id(display_id, msg)
+
+        # set the prompt number for the input and the output
+        if 'execution_count' in content:
+            cell['execution_count'] = content['execution_count']
+
+        if msg_type == 'status':
+            if content['execution_state'] == 'idle':
+                # Set result to None to halt execution
+                result = None
+        elif msg_type == 'clear_output':
+            cell.outputs[:] = []
+            # clear display_id mapping for this cell
+            for display_id, cell_map in self._display_id_map.items():
+                if cell_index in cell_map:
+                    cell_map[cell_index] = []
+        # Check for remaining messages we don't process
+        elif not (msg_type in ['execute_input', 'update_display_data'] or msg_type.startswith('comm')):
             try:
-                out = output_from_msg(msg)
+                # Assign output as our processed "result"
+                result = output_from_msg(msg)
             except ValueError:
                 self.log.error("unhandled iopub msg: " + msg_type)
-                continue
-            if display_id:
-                # record output index in:
-                #   _display_id_map[display_id][cell_idx]
-                cell_map = self._display_id_map.setdefault(display_id, {})
-                output_idx_list = cell_map.setdefault(cell_index, [])
-                output_idx_list.append(len(outs))
+            else:
+                if display_id:
+                    # record output index in:
+                    #   _display_id_map[display_id][cell_idx]
+                    cell_map = self._display_id_map.setdefault(display_id, {})
+                    output_idx_list = cell_map.setdefault(cell_index, [])
+                    output_idx_list.append(len(cell.outputs))
 
-            outs.append(out)
+                cell.outputs.append(result)
 
-        return exec_reply, outs
+        return result
 
 
 def executenb(nb, cwd=None, km=None, **kwargs):