Jelajahi Sumber

DatePeriod: added support for pyyaml

Fabian Peter Hammerle 9 tahun lalu
induk
melakukan
dc19d1827e
2 mengubah file dengan 98 tambahan dan 0 penghapusan
  1. 28 0
      ioex/__init__.py
  2. 70 0
      tests/test_dateperiod_yaml.py

+ 28 - 0
ioex/__init__.py

@@ -6,6 +6,10 @@ import os
 import re
 import sys
 import threading
+try:
+    import yaml
+except ImportError:
+    yaml = None
 
 class UnsupportedLocaleSettingError(locale.Error):
     pass
@@ -51,6 +55,8 @@ def int_input_with_default(prompt, default):
 
 class DatePeriod(object):
 
+    yaml_tag = u'!period'
+
     _timestamp_iso_format = r'\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(.\d+)?(Z|[-\+]\d{2}:\d{2})?'
     _timeperiod_iso_format = r'(?P<start>%(t)s)\/(?P<end>%(t)s)' % {'t': _timestamp_iso_format}
 
@@ -110,3 +116,25 @@ class DatePeriod(object):
         return (type(self) == type(other)
             and self.start == other.start
             and self.end == other.end)
+
+    @classmethod
+    def from_yaml(cls, loader, node):
+        return cls(**loader.construct_mapping(node))
+
+    @classmethod
+    def to_yaml(cls, dumper, period):
+        return dumper.represent_mapping(
+            tag = cls.yaml_tag,
+            mapping = {
+                'start': period.start,
+                'end': period.end,
+                },
+            # represent datetime objects with !timestamp tag
+            flow_style = False,
+            )
+
+if yaml:
+    yaml.add_representer(DatePeriod, DatePeriod.to_yaml)
+    yaml.SafeDumper.add_representer(DatePeriod, DatePeriod.to_yaml)
+    yaml.add_constructor(u'!period', DatePeriod.from_yaml)
+    yaml.SafeLoader.add_constructor(u'!period', DatePeriod.from_yaml)

+ 70 - 0
tests/test_dateperiod_yaml.py

@@ -0,0 +1,70 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+import pytz
+import ioex
+import datetime
+yaml = pytest.importorskip('yaml')
+
+@pytest.mark.parametrize(('period', 'yaml_string'), [
+    [
+        ioex.DatePeriod(
+            start = datetime.datetime(2016, 7, 24, 12, 21, 0),
+            end = datetime.datetime(2016, 7, 24, 12, 22, 13),
+            ),
+        '!period\nstart: 2016-07-24T12:21:00\nend: 2016-07-24T12:22:13',
+        ],
+    [
+        ioex.DatePeriod(
+            start = datetime.datetime(2016, 7, 24, 12, 21, 0),
+            end = datetime.datetime(2016, 7, 24, 12, 22, 13),
+            ),
+        '!period\nstart: 2016-07-24 12:21:00\nend: 2016-07-24 12:22:13',
+        ],
+    [
+        ioex.DatePeriod(
+            start = datetime.datetime(2016, 7, 24, 12, 20, 0, microsecond = 25500),
+            end = datetime.datetime(2016, 7, 24, 12, 21, 0, microsecond = 13),
+            ),
+        '!period\nstart: 2016-07-24T12:20:00.025500\nend: 2016-07-24T12:21:00.000013',
+        ],
+    [
+        ioex.DatePeriod(
+            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),
+            ),
+        '!period\nstart: 2016-07-24T12:20:00.025500\nend: 2016-07-24T12:21:00.000013Z',
+        ],
+    ])
+def test_dateperiod_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)
+
+@pytest.mark.parametrize(('period', 'yaml_string'), [
+    [
+        ioex.DatePeriod(
+            start = datetime.datetime(2016, 7, 24, 12, 21, 0),
+            end = datetime.datetime(2016, 7, 24, 12, 22, 13),
+            ),
+        '!period\nend: 2016-07-24 12:22:13\nstart: 2016-07-24 12:21:00\n',
+        ],
+    [
+        ioex.DatePeriod(
+            start = datetime.datetime(2016, 7, 24, 12, 20, 0, microsecond = 25500),
+            end = datetime.datetime(2016, 7, 24, 12, 21, 0, microsecond = 13),
+            ),
+        '!period\nend: 2016-07-24 12:21:00.000013\nstart: 2016-07-24 12:20:00.025500\n',
+        ],
+    [
+        ioex.DatePeriod(
+            start = pytz.timezone('Europe/London').localize(datetime.datetime(2016, 7, 24, 12, 20, 0)),
+            end = pytz.utc.localize(datetime.datetime(2016, 7, 24, 12, 21, 0, microsecond = 13)),
+            ),
+        '!period\nend: 2016-07-24 12:21:00.000013+00:00\nstart: 2016-07-24 12:20:00+01:00\n',
+        ],
+    ])
+def test_dateperiod_to_yaml(period, yaml_string):
+    assert yaml.dump(period) == yaml_string
+    assert yaml.safe_dump(period) == yaml_string