[ExecuteTime] fix timezone handling using notebook funtions

ipython-contrib · Oct 27, 2017 · a1baaff · a1baaff
1 parent 117b523
commit a1baaff
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 9 deletions.
diff --git a/src/jupyter_contrib_nbextensions/nbconvert_support/execute_time.py b/src/jupyter_contrib_nbextensions/nbconvert_support/execute_time.py
@@ -5,6 +5,7 @@
 from datetime import datetime
 
 from nbconvert.preprocessors.execute import ExecutePreprocessor
+from notebook.services.contents import tz as nbtz
 
 
 class ExecuteTimePreprocessor(ExecutePreprocessor):
@@ -20,13 +21,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
diff --git a/tests/test_preprocessors.py b/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')