Merge pull request #65 from tomalrussell/feature/fast_yaml

Feature/fast yaml [Delivers #142861181]

smif/data_layer/load.py

@@ -2,6 +2,12 @@
 """Parse yaml config files, to construct sector models
 """
 import yaml
+from smif import SpaceTimeValue
+
+try:
+    from yaml import CLoader as Loader, CDumper as Dumper
+except ImportError:
+    from yaml import Loader, Dumper
 
 
 def load(file_path):
@@ -13,7 +19,7 @@ def load(file_path):
         The path of the configuration file to parse
     """
     with open(file_path, 'r') as file_handle:
-        return yaml.load(file_handle)
+        return yaml.load(file_handle, Loader=Loader)
 
 
 def dump(data, file_path):
@@ -27,4 +33,35 @@ def dump(data, file_path):
         Data to write (should be lists, dicts and simple values)
     """
     with open(file_path, 'w') as file_handle:
-        return yaml.dump(data, file_handle)
+        return yaml.dump(data, file_handle, Dumper=Dumper)
+
+
+def space_time_value_representer(dumper, data):
+    """Dump custom yaml representation of SpaceTimeValue
+    """
+    return dumper.represent_sequence(
+        "SpaceTimeValue", [
+            data.region,
+            data.interval,
+            data.value,
+            data.units
+        ]
+    )
+
+
+yaml.add_representer(SpaceTimeValue, space_time_value_representer, Dumper=Dumper)
+
+
+def space_time_value_constructor(loader, node):
+    """Load ustom yaml representation of SpaceTimeValue
+    """
+    value = loader.construct_sequence(node)
+    return SpaceTimeValue(
+        value[0],
+        value[1],
+        value[2],
+        value[3]
+    )
+
+
+yaml.add_constructor("SpaceTimeValue", space_time_value_constructor, Loader=Loader)

tests/data_layer/test_load.py

@@ -0,0 +1,23 @@
+"""Test custom pieces of YAML input/output
+"""
+from smif import SpaceTimeValue
+from smif.data_layer.load import dump, load
+
+
+def test_dump_space_time_value(setup_folder_structure):
+    """Test custom YAML representation
+    """
+    test_file = str(setup_folder_structure.join("test.yaml"))
+    data = SpaceTimeValue("Middlesborough", 0, 33, "GW")
+    dump(data, test_file)
+    with open(test_file, 'r') as fh:
+        line = fh.readline()
+        assert line.strip() == "!<SpaceTimeValue> [Middlesborough, 0, 33, GW]"
+
+
+def test_load_space_time_value(setup_folder_structure):
+    test_file = str(setup_folder_structure.join("test.yaml"))
+    with open(test_file, 'w') as fh:
+        fh.write("!<SpaceTimeValue> [Middlesborough, 0, 33, GW]")
+    data = load(test_file)
+    assert data == SpaceTimeValue("Middlesborough", 0, 33, "GW")