calcex.Figure.register_yaml_constructor(): added param 'tag'

ioex/calcex.py

@@ -77,7 +77,10 @@ class Figure(object):
     @classmethod
     def to_yaml(cls, dumper, figure, tag=yaml_tag):
         if figure.value is None or type(figure.value) in [int, float]:
-            value_text = u'?' if figure.value is None else u'{}'.format(figure.value)
+            if figure.value is None:
+                value_text = u'?'
+            else:
+                value_text = u'{}'.format(figure.value)
             if figure.unit is None:
                 figure_text = value_text
             else:
@@ -90,8 +93,8 @@ class Figure(object):
             )
 
     @classmethod
-    def register_yaml_representer(cls, dumper):
-        dumper.add_representer(cls, cls.to_yaml)
+    def register_yaml_representer(cls, dumper, tag=yaml_tag):
+        dumper.add_representer(cls, lambda d, f: cls.to_yaml(d, f, tag=tag))
 
     def __eq__(self, other):
         return isinstance(self, type(other)) and vars(self) == vars(other)

tests/calcex/test_figure_yaml.py

@@ -95,11 +95,11 @@ def test_to_yaml_mapping(figure, yaml_dumper):
 def test_register_yaml_representer(figure, yaml_dumper):
     class TestDumper(yaml_dumper):
         pass
-    figure.register_yaml_representer(TestDumper)
+    figure.register_yaml_representer(TestDumper, tag='!test-fig')
     generated_yaml = yaml.dump(figure, Dumper=TestDumper)
 
     class TestLoader(yaml.SafeLoader):
         pass
-    figure.register_yaml_constructor(TestLoader, tag='!figure')
+    figure.register_yaml_constructor(TestLoader, tag='!test-fig')
     loaded_figure = yaml.load(generated_yaml, Loader=TestLoader)
     assert figure == loaded_figure