[ExecuteTime] fix timezone handling using notebook funtions

Author: jcb91

src/jupyter_contrib_nbextensions/nbconvert_support/execute_time.py

@@ -6,6 +6,13 @@
 
 from nbconvert.preprocessors.execute import ExecutePreprocessor
 
+try:
+    # notebook >= 5.0.0-rc1
+    import notebook._tz as nbtz
+except ImportError:
+    # notebook < 5.0.0-rc1
+    import notebook.services.contents.tz as nbtz
+
 
 class ExecuteTimePreprocessor(ExecutePreprocessor):
     """
@@ -20,13 +27,13 @@ def run_cell(self, cell, cell_index, *args, **kwargs):
         if exec_reply.get('msg_type', '') == 'execute_reply':
             ets = cell.setdefault('metadata', {}).setdefault('ExecuteTime', {})
             if 'started' in exec_reply.get('metadata', {}):
-                # started value should is already a string
+                # started value should is already a string, so don't isoformat
                 ets['start_time'] = exec_reply['metadata']['started']
             else:
                 # attempt to fallback to datetime obj for execution request msg
-                ets['start_time'] = exec_reply.get('parent_header', {}).get(
-                    'date', before).isoformat()
-            ets['end_time'] = (exec_reply.get('header', {}).get(
-                'date', None) or datetime.utcnow()).isoformat()
+                ets['start_time'] = nbtz.isoformat(
+                    exec_reply.get('parent_header', {}).get('date', before))
+            ets['end_time'] = nbtz.isoformat(
+                exec_reply.get('header', {}).get('date') or nbtz.utcnow())
 
         return exec_reply, outs

tests/test_preprocessors.py

@@ -1,8 +1,8 @@
 # -*- coding: utf-8 -*-
 
-import datetime
 import json
 import os
+import re
 
 import nbformat.v4 as nbf
 from nbconvert import LatexExporter, NotebookExporter, RSTExporter
@@ -113,6 +113,16 @@ def test_preprocessor_svg2pdf():
               'exported pdf should be referenced in exported notebook')
 
 
+def _normalize_iso8601_timezone(timestamp_str):
+    # Zulu -> +00:00 offset
+    timestamp_str = re.sub(r'Z$', r'+00:00', timestamp_str)
+    # HH -> HH:00 offset
+    timestamp_str = re.sub(r'([+-]\d\d)$', r'\1:00', timestamp_str)
+    # HHMM -> HH:MM offset
+    timestamp_str = re.sub(r'([+-]\d\d):?(\d\d)$', r'\1:\2', timestamp_str)
+    return timestamp_str
+
+
 def test_preprocessor_execute_time():
     """Test ExecuteTime preprocessor."""
     # check import shortcut
@@ -133,6 +143,6 @@ def test_preprocessor_execute_time():
             assert_in('start_time', etmd)
             assert_in('end_time', etmd)
             assert_greater_equal(
-                datetime.datetime.utcfromtimestamp(etmd['end_time']),
-                datetime.datetime.utcfromtimestamp(etmd['start_time']),
-                'end_time should be after start time')
+                _normalize_iso8601_timezone(etmd['end_time']),
+                _normalize_iso8601_timezone(etmd['start_time']),
+                'end_time should not be before start_time')