Bläddra i källkod

calcex.Figure: added register_yaml_constructor()

Fabian Peter Hammerle 8 år sedan
förälder
incheckning
e1ede65920
2 ändrade filer med 23 tillägg och 4 borttagningar
  1. 12 4
      ioex/calcex.py
  2. 11 0
      tests/calcex/test_figure_yaml.py

+ 12 - 4
ioex/calcex.py

@@ -49,11 +49,19 @@ class Figure(object):
             return '{} {}'.format(self.value, self.unit)
 
     @classmethod
-    def to_yaml(cls, dumper, figure, tag = yaml_tag):
+    def from_yaml(cls, loader, node):
+        return cls(**loader.construct_mapping(node))
+
+    @classmethod
+    def register_yaml_constructor(cls, loader, tag=yaml_tag):
+        loader.add_constructor(tag, cls.from_yaml)
+
+    @classmethod
+    def to_yaml(cls, dumper, figure, tag=yaml_tag):
         return dumper.represent_mapping(
-            tag = tag,
-            mapping = {'value': figure.value, 'unit': figure.unit},
-            )
+            tag=tag,
+            mapping={'value': figure.value, 'unit': figure.unit},
+        )
 
     @classmethod
     def register_yaml_representer(cls, dumper):

+ 11 - 0
tests/calcex/test_figure_yaml.py

@@ -4,6 +4,17 @@ import pytest
 from ioex.calcex import Figure, UnitMismatchError
 yaml = pytest.importorskip('yaml')
 
+@pytest.mark.parametrize(('yaml_loader'), [yaml.Loader, yaml.SafeLoader])
+@pytest.mark.parametrize(('figure_yaml', 'expected_figure'), [
+    ['!fig {value: null, unit: null}', Figure()],
+    ['!fig {value: 123.4, unit: null}', Figure(123.4)],
+    ['!fig {value: [1, 2], unit: null}', Figure([1, 2])],
+])
+def test_register_yaml_constructor(figure_yaml, expected_figure, yaml_loader):
+    class TestLoader(yaml_loader):
+        pass
+    Figure.register_yaml_constructor(TestLoader, tag = '!fig')
+    assert expected_figure == yaml.load(figure_yaml, Loader = TestLoader)
 
 @pytest.mark.parametrize(('yaml_dumper'), [yaml.Dumper, yaml.SafeDumper])
 @pytest.mark.parametrize(('figure'), [