derive figure from yamlobject

dingguo/__init__.py

@@ -1,9 +1,14 @@
 # -*- coding: utf-8 -*-
 
 import yaml
+import yaml.representer
 import datetime
 
-class Figure(object):
+yaml.Dumper.add_representer(unicode, yaml.representer.SafeRepresenter.represent_unicode)
+
+class Figure(yaml.YAMLObject):
+
+    yaml_tag = u"!figure"
 
     def __init__(self, value, unit):
         self.value = value
@@ -34,6 +39,22 @@ class Figure(object):
     def __ne__(self, other):
         return not (self == other)
 
+    @classmethod
+    def from_yaml(cls, loader, node):
+        attr = loader.construct_mapping(node)
+        print(type(attr['unit']), attr['unit'])
+        return cls(
+                value = attr['value'],
+                unit = unicode(attr['unit']),
+                )
+
+    @classmethod
+    def to_yaml(cls, dumper, figure):
+        return dumper.represent_mapping(
+                cls.yaml_tag,
+                {'unit': figure.get_unit(), 'value': figure.get_value()},
+                )
+
 class Distance(Figure):
 
     def __init__(self, value, unit):

tests/test_.py

@@ -7,6 +7,9 @@ import dingguo
 import yaml
 import os
 
+def test_sum_init():
+    assert dingguo.Sum(1.23, u'EUR') == dingguo.Sum(1.23, u'€')
+
 def test_order_dict_repr():
 
     order = dingguo.Order(

tests/test_yaml.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+
+import pytest
+
+import os
+import yaml
+import dingguo
+import datetime
+
+project_root_path = os.path.realpath(os.path.join(__file__, '..', '..'))
+test_data_path = os.path.join(project_root_path, 'tests', 'data')
+
+def get_figure_a():
+    return dingguo.Figure(12.3, u'km')
+
+def get_figure_b():
+    return dingguo.Figure(12300, u'米')
+
+def to_yaml(data):
+    return yaml.dump(data, default_flow_style = False, allow_unicode = True).decode('utf-8')
+
+def test_figure_to_yaml():
+    assert to_yaml(get_figure_a()) == u"""!figure
+unit: km
+value: 12.3
+"""
+
+def test_figure_to_yaml_unicode():
+    assert to_yaml(get_figure_b()) == u"""!figure
+unit: 米
+value: 12300
+"""
+
+def test_figure_from_yaml():
+    assert get_figure_a() == yaml.load(u"""!figure
+unit: km
+value: 12.3
+""")
+
+def test_figure_to_yaml_unicode():
+    assert get_figure_b() == yaml.load(u"""!figure
+unit: 米
+value: 12300
+""")