added add_yaml_timestamp_constructor() and construct_yaml_timestamp() to fix pyyaml not setting tzinfo attribute of datetime instances while loading

ioex/datetimeex.py

@@ -0,0 +1,21 @@
+import datetime
+import dateutil.tz.tz
+import re
+try:
+    import yaml
+except ImportError:
+    yaml = None
+
+def construct_yaml_timestamp(loader, node):
+    loaded_dt = loader.construct_yaml_timestamp(node)
+    if type(loaded_dt) is datetime.datetime and loaded_dt.tzinfo is None:
+        timezone_match = re.search(
+            ur'(Z|[\+-]\d{2}:\d{2})$',
+            loader.construct_yaml_str(node),
+            )
+        if timezone_match:
+            loaded_dt = loaded_dt.replace(tzinfo = dateutil.tz.tz.tzutc())
+    return loaded_dt
+
+def add_yaml_timestamp_constructor(loader, tag = u'tag:yaml.org,2002:timestamp'):
+    loader.add_constructor(tag, construct_yaml_timestamp)

tests/datetimeex/test_yaml_timestamp_constructor.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+yaml = pytest.importorskip('yaml')
+import copy
+import datetime
+import dateutil.tz.tz
+import ioex.datetimeex
+import pytz
+
+@pytest.mark.parametrize(('loader'), [yaml.Loader, yaml.SafeLoader])
+@pytest.mark.parametrize(('yaml_string', 'expected_timestamp'), [
+    ['2016-07-14 13:50:04', datetime.datetime(2016, 7, 14, 13, 50, 4, 0)],
+    ['2016-07-14 13:50:04Z', pytz.utc.localize(datetime.datetime(2016, 7, 14, 13, 50, 4, 0))],
+    ['2016-07-14 13:50:04Z', datetime.datetime(2016, 7, 14, 13, 50, 4, 0, tzinfo = dateutil.tz.tz.tzutc())],
+    ['2016-01-14 13:50:04+01:00', pytz.timezone('Europe/Vienna').localize(datetime.datetime(2016, 1, 14, 13, 50, 4, 0))],
+    ['2016-07-14 13:50:04+02:00', pytz.timezone('Europe/Vienna').localize(datetime.datetime(2016, 7, 14, 13, 50, 4, 0))],
+    ['2016-07-14 13:50:04+02:00', datetime.datetime(2016, 7, 14, 13, 50, 4, 0, tzinfo = dateutil.tz.tz.tzoffset('Vienna', 2 * 60 * 60))],
+    ])
+def test_from_yaml(yaml_string, expected_timestamp, loader):
+    loader_copy = copy.deepcopy(loader)
+    ioex.datetimeex.add_yaml_timestamp_constructor(loader_copy)
+    assert yaml.load(yaml_string, Loader = loader_copy) == expected_timestamp
+
+@pytest.mark.parametrize(('yaml_string', 'tag', 'expected_timestamp'), [
+    ['!without_timezone 2016-07-14 13:50:04', '!without_timezone', datetime.datetime(2016, 7, 14, 13, 50, 4, 0)],
+    ['!datetime 2016-07-14 13:50:04Z', '!datetime', pytz.utc.localize(datetime.datetime(2016, 7, 14, 13, 50, 4, 0))],
+    ])
+def test_from_yaml_tag(yaml_string, tag, expected_timestamp):
+    loader = copy.deepcopy(yaml.SafeLoader)
+    ioex.datetimeex.add_yaml_timestamp_constructor(loader, tag = tag)
+    assert yaml.load(yaml_string, Loader = loader) == expected_timestamp

tests/test_yaml.py → tests/external/test_yaml.py

@@ -1,9 +1,9 @@
 # -*- coding: utf-8 -*-
 import pytest
 
-import pytz
-import datetime
 yaml = pytest.importorskip('yaml')
+import datetime
+import pytz
 
 @pytest.mark.parametrize('source_object,yaml_string', [
     [datetime.datetime(2016, 7, 14, 13, 50, 4, 0), '2016-07-14 13:50:04\n...\n'],