calcex.Figure: added register_yaml_representer()

ioex/calcex.py

@@ -1,8 +1,8 @@
 import copy
-# try:
-#     import yaml
-# except ImportError:
-#     yaml = None
+try:
+    import yaml
+except ImportError:
+    yaml = None
 
 
 class UnitMismatchError(ValueError):
@@ -48,6 +48,17 @@ class Figure(object):
         else:
             return '{} {}'.format(self.value, self.unit)
 
+    @classmethod
+    def to_yaml(cls, dumper, figure, tag = yaml_tag):
+        return dumper.represent_mapping(
+            tag = tag,
+            mapping = {'value': figure.value, 'unit': figure.unit},
+            )
+
+    @classmethod
+    def register_yaml_representer(cls, dumper):
+        dumper.add_representer(cls, cls.to_yaml)
+
     def __eq__(self, other):
         return isinstance(self, type(other)) and vars(self) == vars(other)
 

tests/calcex/test_figure.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 import pytest
 
-import copy
 from ioex.calcex import Figure, UnitMismatchError
 
 

tests/calcex/test_figure_yaml.py

@@ -0,0 +1,54 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+from ioex.calcex import Figure, UnitMismatchError
+yaml = pytest.importorskip('yaml')
+
+
+@pytest.mark.parametrize(('yaml_dumper'), [yaml.Dumper, yaml.SafeDumper])
+@pytest.mark.parametrize(('figure'), [
+    Figure(),
+    Figure(123.4),
+    Figure([1, 2]),
+    Figure(None, u'm/s²'),
+    Figure(123.4, u'm/s²'),
+    Figure(1234, u'米/s²'),
+])
+def test_to_yaml(figure, yaml_dumper):
+    class TestDumper(yaml_dumper):
+        pass
+    TestDumper.add_representer(
+        type(figure),
+        lambda d, f: figure.to_yaml(d, f, '!test-figure'))
+    figure_yaml = yaml.dump(figure, Dumper=TestDumper)
+    class TestLoader(yaml.SafeLoader):
+        pass
+    TestLoader.add_constructor(
+        '!test-figure',
+        lambda loader, node: loader.construct_mapping(node),
+        )
+    figure_attr = yaml.load(figure_yaml, Loader=TestLoader)
+    assert set(['value', 'unit']) == set(figure_attr.keys())
+    assert figure.value == figure_attr['value']
+    assert figure.unit == figure_attr['unit']
+
+
+@pytest.mark.parametrize(('yaml_dumper'), [yaml.Dumper, yaml.SafeDumper])
+@pytest.mark.parametrize(('figure'), [
+    Figure(123.4, u'm/s²'),
+])
+def test_register_yaml_representer(figure, yaml_dumper):
+    class TestDumper(yaml_dumper):
+        pass
+    figure.register_yaml_representer(TestDumper)
+    figure_yaml = yaml.dump(figure, Dumper=TestDumper)
+    class TestLoader(yaml.SafeLoader):
+        pass
+    TestLoader.add_constructor(
+        '!figure',
+        lambda loader, node: loader.construct_mapping(node),
+        )
+    figure_attr = yaml.load(figure_yaml, Loader=TestLoader)
+    assert set(['value', 'unit']) == set(figure_attr.keys())
+    assert figure.value == figure_attr['value']
+    assert figure.unit == figure_attr['unit']