datetimeex.Period: timezone while loading from yaml not set fixed

ioex/datetimeex.py

@@ -110,8 +110,19 @@ class Period(object):
             flow_style = False,
             )
 
+    def __repr__(self):
+        return '%s(%s)' % (self.__class__.__name__, ', '.join([
+            'start = %r' % self.start,
+            'end = %r' % self.end,
+            ]))
+
+    @classmethod
+    def add_yaml_constructor(cls, loader, tag = yaml_tag):
+        add_yaml_timestamp_constructor(loader)
+        loader.add_constructor(tag, cls.from_yaml)
+
 if yaml:
     yaml.add_representer(Period, Period.to_yaml)
     yaml.SafeDumper.add_representer(Period, Period.to_yaml)
-    yaml.add_constructor(u'!period', Period.from_yaml)
-    yaml.SafeLoader.add_constructor(u'!period', Period.from_yaml)
+    Period.add_yaml_constructor(yaml.Loader)
+    Period.add_yaml_constructor(yaml.SafeLoader)

tests/datetimeex/test_period_yaml.py

@@ -1,12 +1,14 @@
 # -*- coding: utf-8 -*-
 import pytest
 
-import pytz
-import ioex.datetimeex
+import copy
 import datetime
+import ioex.datetimeex
+import pytz
 yaml = pytest.importorskip('yaml')
 
-@pytest.mark.parametrize(('period', 'yaml_string'), [
+@pytest.mark.parametrize(('loader'), [yaml.Loader, yaml.SafeLoader])
+@pytest.mark.parametrize(('expected_period', 'yaml_string'), [
     [
         ioex.datetimeex.Period(
             start = datetime.datetime(2016, 7, 24, 12, 21, 0),
@@ -31,16 +33,24 @@ yaml = pytest.importorskip('yaml')
     [
         ioex.datetimeex.Period(
             start = datetime.datetime(2016, 7, 24, 12, 20, 0, microsecond = 25500),
-            end = datetime.datetime(2016, 7, 24, 12, 21, 0, microsecond = 13, tzinfo = pytz.utc),
+            end = pytz.utc.localize(datetime.datetime(2016, 7, 24, 12, 21, 0, microsecond = 13)),
             ),
         '!period\nstart: 2016-07-24T12:20:00.025500\nend: 2016-07-24T12:21:00.000013Z',
         ],
+    [
+        ioex.datetimeex.Period(
+            start = pytz.timezone('Europe/London').localize(datetime.datetime(2016, 1, 24, 12, 20, 0, microsecond = 25500)),
+            end = pytz.timezone('Europe/London').localize(datetime.datetime(2016, 7, 24, 12, 21, 0, microsecond = 13)),
+            ),
+        '!period\nstart: 2016-01-24T12:20:00.025500Z\nend: 2016-07-24T12:21:00.000013+01:00',
+        ],
     ])
-def test_from_yaml(period, yaml_string):
-    if period.start.tzinfo or period.end.tzinfo:
-        pytest.xfail('pyyaml ignores timezones when loading timestamps')
-    assert period == yaml.load(yaml_string)
-    assert period == yaml.safe_load(yaml_string)
+def test_from_yaml(expected_period, yaml_string, loader):
+    loader_copy = copy.deepcopy(loader)
+    loaded_period = yaml.load(yaml_string, Loader = loader_copy)
+    assert expected_period == loaded_period
+    assert expected_period.start.utcoffset() == loaded_period.start.utcoffset()
+    assert expected_period.end.utcoffset() == loaded_period.end.utcoffset()
 
 @pytest.mark.parametrize(('period', 'yaml_string'), [
     [

tests/datetimeex/test_yaml_timestamp_constructor.py

@@ -33,3 +33,14 @@ 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
+
+@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))],
+    ])
+def test_from_yaml_repeat(yaml_string, expected_timestamp):
+    loader = copy.deepcopy(yaml.SafeLoader)
+    ioex.datetimeex.add_yaml_timestamp_constructor(loader)
+    assert yaml.load(yaml_string, Loader = loader) == expected_timestamp
+    ioex.datetimeex.add_yaml_timestamp_constructor(loader)
+    assert yaml.load(yaml_string, Loader = loader) == expected_timestamp