Safe yaml loading and dumping (#5)

* YAML safe load and dump
* SerializationSafetyError and deprecation warnings
* Unit tests

baiji/serialization/__init__.py

@@ -7,4 +7,4 @@
 from pkgutil import extend_path
 __path__ = extend_path(__path__, __name__)
 
-__version__ = '1.0.2'
+__version__ = '1.1.0'

baiji/serialization/test_yaml.py

@@ -0,0 +1,26 @@
+import unittest
+from baiji.serialization import yaml
+
+class UnsafeToDump(object):
+    """Python classes should not by default be YAML dumpable in safe mode"""
+    pass
+
+unsafe_to_load = "!!python/name:baiji.serialization.test_yaml.UnsafeToDump ''\n"
+
+class TestYAML(unittest.TestCase):
+
+    def test_yaml_unsafe_dumps(self):
+        self.assertEqual(yaml.dumps(UnsafeToDump), unsafe_to_load)
+
+    def test_yaml_unsafe_loads_uses_load(self):
+        self.assertEqual(yaml.loads(unsafe_to_load), UnsafeToDump)
+
+    def test_yaml_safe_dump_raises_on_unsafe(self):
+        self.assertRaises(
+            yaml.SerializationSafetyError,
+            yaml.dumps, UnsafeToDump, safe=True)
+
+    def test_yaml_safe_load_raises_on_unsafe(self):
+        self.assertRaises(
+            yaml.SerializationSafetyError,
+            yaml.loads, unsafe_to_load, safe=True)

baiji/serialization/yaml.py

@@ -2,6 +2,8 @@
 
 EXTENSION = '.yaml'
 
+class SerializationSafetyError(Exception):
+    pass
 
 def dump(obj, f, *args, **kwargs):
     from baiji.serialization.util.openlib import ensure_file_open_and_call
@@ -14,20 +16,62 @@ def load(f, *args, **kwargs):
 
 
 def loads(s, *args, **kwargs):
-    import yaml
-    return yaml.load(s, *args, **kwargs)
+    return _load(s, *args, **kwargs)
 
 
 def dumps(obj, *args, **kwargs):
     import yaml
-    return yaml.dump(obj, *args, **kwargs)
+    import warnings
+    safe_mode = kwargs.pop('safe', None)
+    try:
+        return yaml.safe_dump(obj, *args, **kwargs)
+    except yaml.representer.RepresenterError as e:
+        if safe_mode:
+            raise SerializationSafetyError(*e.args)
+        if safe_mode is None:
+            warnings.warn(
+                ('Unsafe YAML serialization. This will generate an error in the '
+                 'future. Call with `safe=False` if you really want unsafe serialization.'),
+                DeprecationWarning, stacklevel=2
+            )
+        # if safe_mode is None (not given) or False (explicitly set), fall back to unsafe behavior
+        return yaml.dump(obj, *args, **kwargs)
+
 
 
 def _dump(f, obj, *args, **kwargs):
     import yaml
-    return yaml.dump(obj, f, *args, **kwargs)
+    import warnings
+    safe_mode = kwargs.pop('safe', None)
+    try:
+        return yaml.safe_dump(obj, f, *args, **kwargs)
+    except yaml.representer.RepresenterError as e:
+        if safe_mode:
+            raise SerializationSafetyError(*e.args)
+        if safe_mode is None:
+            warnings.warn(
+                ('Unsafe YAML serialization. This will generate an error in the '
+                 'future. Call with `safe=False` if you really want unsafe serialization.'),
+                DeprecationWarning, stacklevel=2
+            )
+        # if safe_mode is None (not given) or False (explicitly set), fall back to unsafe behavior
+        return yaml.dump(obj, f, *args, **kwargs)
 
 
 def _load(f, *args, **kwargs):
     import yaml
-    return yaml.load(f, *args, **kwargs)
+    import warnings
+    safe_mode = kwargs.pop('safe', None)
+    try:
+        return yaml.safe_load(f, *args, **kwargs)
+    except yaml.constructor.ConstructorError as e:
+        if safe_mode:
+            raise SerializationSafetyError(*e.args)
+        if safe_mode is None:
+            warnings.warn(
+                ('Unsafe YAML serialization. This will generate an error in the '
+                 'future. Call with `safe=False` if you really want unsafe serialization.'),
+                DeprecationWarning, stacklevel=2
+            )
+        # if safe_mode is None (not given) or False (explicitly set), fall back to unsafe behavior
+        return yaml.load(f, *args, **kwargs)