Add execution timing in code cell metadata for v4 spec (#32)

* Add execution timing in code cell metadata for v4 spec

* Add record_timing trait

* Add test

Author: davidbrochart

nbclient/client.py

@@ -1,3 +1,4 @@
+import datetime
 import base64
 from textwrap import dedent
 
@@ -17,6 +18,10 @@
 from .exceptions import CellTimeoutError, DeadKernelError, CellExecutionComplete, CellExecutionError
 
 
+def timestamp():
+    return datetime.datetime.utcnow().isoformat() + 'Z'
+
+
 class NotebookClient(LoggingConfigurable):
     """
     Encompasses a Client for executing cells in a notebook
@@ -144,6 +149,16 @@ class NotebookClient(LoggingConfigurable):
         ),
     ).tag(config=True)
 
+    record_timing = Bool(
+        True,
+        help=dedent(
+            """
+            If `True` (default), then the execution timings of each cell will
+            be stored in the metadata of the notebook.
+            """
+        ),
+    ).tag(config=True)
+
     iopub_timeout = Integer(
         4,
         allow_none=False,
@@ -436,6 +451,8 @@ async def _poll_for_reply(self, msg_id, cell, timeout, task_poll_output_msg):
             try:
                 msg = await self.kc.shell_channel.get_msg(timeout=timeout)
                 if msg['parent_header'].get('msg_id') == msg_id:
+                    if self.record_timing:
+                        cell['metadata']['execution']['shell.execute_reply'] = timestamp()
                     try:
                         await asyncio.wait_for(task_poll_output_msg, self.iopub_timeout)
                     except (asyncio.TimeoutError, Empty):
@@ -608,6 +625,9 @@ async def async_execute_cell(self, cell, cell_index, execution_count=None, store
             self.log.debug("Skipping non-executing cell %s", cell_index)
             return cell
 
+        if self.record_timing and 'execution' not in cell['metadata']:
+            cell['metadata']['execution'] = {}
+
         self.log.debug("Executing cell:\n%s", cell.source)
         parent_msg_id = self.kc.execute(
             cell.source, store_history=store_history, stop_on_error=not self.allow_errors
@@ -672,6 +692,15 @@ def process_message(self, msg, cell, cell_index):
         if 'execution_count' in content:
             cell['execution_count'] = content['execution_count']
 
+        if self.record_timing:
+            if msg_type == 'status':
+                if content['execution_state'] == 'idle':
+                    cell['metadata']['execution']['iopub.status.idle'] = timestamp()
+                elif content['execution_state'] == 'busy':
+                    cell['metadata']['execution']['iopub.status.busy'] = timestamp()
+            elif msg_type == 'execute_input':
+                cell['metadata']['execution']['iopub.execute_input'] = timestamp()
+
         if msg_type == 'status':
             if content['execution_state'] == 'idle':
                 raise CellExecutionComplete()

nbclient/tests/files/Sleep1s.ipynb

@@ -0,0 +1,65 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "import datetime"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "t0 = datetime.datetime.utcnow()\n",
+    "time.sleep(1)\n",
+    "t1 = datetime.datetime.utcnow()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "time_format = '%Y-%m-%dT%H:%M:%S.%fZ'\n",
+    "print(t0.strftime(time_format), end='')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(t1.strftime(time_format), end='')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

nbclient/tests/test_client.py

@@ -5,6 +5,7 @@
 import re
 import threading
 import asyncio
+import datetime
 
 import nbformat
 import sys
@@ -354,6 +355,36 @@ def test_many_async_parallel_notebooks(capfd):
     assert captured.err == ""
 
 
+def test_execution_timing():
+    """Compare the execution timing information stored in the cell with the
+    actual time it took to run the cell. Also check for the cell timing string
+    format."""
+    opts = dict(kernel_name="python")
+    input_name = "Sleep1s.ipynb"
+    input_file = os.path.join(current_dir, "files", input_name)
+    res = notebook_resources()
+    input_nb, output_nb = run_notebook(input_file, opts, res)
+
+    def get_time_from_str(s):
+        time_format = '%Y-%m-%dT%H:%M:%S.%fZ'
+        return datetime.datetime.strptime(s, time_format)
+
+    execution_timing = output_nb['cells'][1]['metadata']['execution']
+    status_busy = get_time_from_str(execution_timing['iopub.status.busy'])
+    execute_input = get_time_from_str(execution_timing['iopub.execute_input'])
+    execute_reply = get_time_from_str(execution_timing['shell.execute_reply'])
+    status_idle = get_time_from_str(execution_timing['iopub.status.idle'])
+
+    cell_start = get_time_from_str(output_nb['cells'][2]['outputs'][0]['text'])
+    cell_end = get_time_from_str(output_nb['cells'][3]['outputs'][0]['text'])
+
+    delta = datetime.timedelta(milliseconds=100)
+    assert status_busy - cell_start < delta
+    assert execute_input - cell_start < delta
+    assert execute_reply - cell_end < delta
+    assert status_idle - cell_end < delta
+
+
 class TestExecute(NBClientTestsBase):
     """Contains test functions for execute.py"""