[io] fix SafeIncludeLoader

flavio/io/test_yaml.py

@@ -9,3 +9,15 @@ def test_load_include(self):
             d = flavio.io.yaml.load_include(f)
         self.assertDictEqual(d, {'key': {'key2': 0},
                                  'key3': [1, 2, 3, 4, 5, 6, 7, 8]})
+    def test_load_include_error(self):
+        test_str = "key: !include relative_path.yaml"
+        with self.assertRaises(ValueError):
+            d = flavio.io.yaml.load_include(test_str)
+        test_str = "key: !include /absolute_path.yaml"
+        with self.assertRaises(OSError):
+            d = flavio.io.yaml.load_include(test_str)
+    def test_include_absolute_path(self):
+        test_file = os.path.join(os.path.dirname(__file__), '..', 'data', 'test', '2.yaml')
+        test_str = "key: !include {}".format(test_file)
+        d = flavio.io.yaml.load_include(test_str)
+        self.assertDictEqual(d, {'key': {'key2': 0}})

flavio/io/yaml.py

@@ -3,6 +3,8 @@
 import yaml
 from collections import OrderedDict
 import os
+import warnings
+import errno
 
 
 def represent_dict_order(self, data):
@@ -20,29 +22,56 @@ class SafeIncludeLoader(yaml.SafeLoader):
     constructors."""
     def __init__(self, stream):
         try:
-            self._root = os.path.split(stream.name)[0]
+            self._root = os.path.dirname(stream.name)
         except AttributeError:
             # this happens when `stream` is a string
             self._root = None
         super().__init__(stream)
 
     def include(self, node):
-        filename = os.path.join(self._root, self.construct_scalar(node))
+        file = self.construct_scalar(node)
+        filename = self._get_filename(file)
+        if filename is None and self._root is None:
+            raise ValueError("'{}' cannot be included. If the YAML input is a string, '!include' only supports absolute paths.".format(file))
         with open(filename, 'r') as f:
             return yaml.load(f, self.__class__)
 
     def include_merge_list(self, node):
+        if self._root is None:
+            warnings.warn("If the YAML input is a string, '!include_merge_list' only supports absolute paths.")
         files = self.construct_sequence(node)
         a = []
         for file in files:
-            try:
-                filename = os.path.join(self._root, file)
-                with open(filename, 'r') as f:
-                    a += yaml.load(f, self.__class__)
-            except (FileNotFoundError, TypeError):
+            filename = self._get_filename(file)
+            if filename is None:
                 a.append(file)
+            else:
+                try:
+                    with open(filename, 'r') as f:
+                        a += yaml.load(f, self.__class__)
+                except OSError as e:
+                    if e.errno in [errno.ENOENT, errno.EINVAL, errno.EISDIR]:
+                        # ENOENT: FileNotFoundError
+                        # EINVAL: Invalid argument
+                        # EISDIR: IsADirectoryError
+                        a.append(file)
+                    else:
+                        raise
         return a
 
+    def _get_filename(self, file):
+        try:
+            if os.path.isabs(file):
+                return file
+            elif self._root is not None:
+                return os.path.join(self._root, file)
+            else:
+                return None
+        except TypeError:
+            return None
+
+
+
 
 SafeIncludeLoader.add_constructor('!include',
                                   SafeIncludeLoader.include)