Expose some information about notebook execution state

Author: Golf Player

nbclient/client.py

@@ -1,6 +1,7 @@
 import atexit
 import collections
 import datetime
+import enum
 import base64
 import signal
 from textwrap import dedent
@@ -32,6 +33,15 @@ def timestamp():
     return datetime.datetime.utcnow().isoformat() + 'Z'
 
 
+class ExecutionState(enum.Enum):
+    NOTHING = 0
+    STARTUP = 1
+    STARTING_KERNEL_CLIENT = 2
+    EXECUTING = 3
+    CLEANING_UP = 4
+    COMPLETE = 5
+
+
 class NotebookClient(LoggingConfigurable):
     """
     Encompasses a Client for executing cells in a notebook
@@ -310,6 +320,9 @@ def reset_execution_trackers(self):
         self.output_hook_stack = collections.defaultdict(list)
         # our front-end mimicing Output widgets
         self.comm_objects = {}
+        self.state = ExecutionState.NOTHING
+        self.current_cell = None
+        self.current_cell_index = -1
 
     def start_kernel_manager(self):
         """Creates a new kernel manager.
@@ -332,6 +345,7 @@ def start_kernel_manager(self):
         return self.km
 
     async def _async_cleanup_kernel(self):
+        self.state = ExecutionState.CLEANING_UP
         now = self.shutdown_kernel == "immediate"
         try:
             # Queue the manager to kill the process, and recover gracefully if it's already dead.
@@ -367,6 +381,8 @@ async def async_start_new_kernel_client(self, **kwargs):
         kernel_id : string-ized version 4 uuid
            The id of the started kernel.
         """
+
+        self.state = ExecutionState.STARTING_KERNEL_CLIENT
         resource_path = self.resources.get('metadata', {}).get('path') or None
         if resource_path and 'cwd' not in kwargs:
             kwargs["cwd"] = resource_path
@@ -499,8 +515,10 @@ async def async_execute(self, reset_kc=False, **kwargs):
         if reset_kc and self.km:
             await self._async_cleanup_kernel()
         self.reset_execution_trackers()
+        self.state = ExecutionState.STARTUP
 
         async with self.async_setup_kernel(**kwargs):
+            self.state = ExecutionState.EXECUTING
             self.log.info("Executing notebook with kernel: %s" % self.kernel_name)
             for index, cell in enumerate(self.nb.cells):
                 # Ignore `'execution_count' in content` as it's always 1
@@ -512,7 +530,7 @@ async def async_execute(self, reset_kc=False, **kwargs):
             info_msg = await self.async_wait_for_reply(msg_id)
             self.nb.metadata['language_info'] = info_msg['content']['language_info']
             self.set_widgets_metadata()
-
+        self.state = ExecutionState.COMPLETE
         return self.nb
 
     execute = run_sync(async_execute)
@@ -711,6 +729,9 @@ async def async_execute_cell(self, cell, cell_index, execution_count=None, store
         if self.record_timing and 'execution' not in cell['metadata']:
             cell['metadata']['execution'] = {}
 
+        self.current_cell = cell
+        self.current_cell_index = cell_index
+
         self.log.debug("Executing cell:\n%s", cell.source)
         parent_msg_id = await ensure_async(
             self.kc.execute(