Add basic support for YAML data files

Add support for YAML files, when PyYAML is installed. If PyYAML is not
installed, any YAML file_data functions error out.

Only files ending with .yml or .yaml are considered to be YAML files.
All other files are loaded as JSON files.

ddt.py

@@ -11,6 +11,13 @@ import os
 import re
 from functools import wraps
 
+try:
+    import yaml
+except Exception:
+    _have_yaml = False
+else:
+    _have_yaml = True
+
 __version__ = '1.0.3'
 
 # These attributes will not conflict with any real python attribute
@@ -152,25 +159,41 @@ def process_file_data(cls, name, func, file_attr):
     cls_path = os.path.abspath(inspect.getsourcefile(cls))
     data_file_path = os.path.join(os.path.dirname(cls_path), file_attr)
 
-    def _raise_ve(*args):  # pylint: disable-msg=W0613
+    def _raise_not_exists_error(*args):  # pylint: disable-msg=W0613
         raise ValueError("%s does not exist" % file_attr)
 
+    def _raise_need_yaml_error(*args):  # pylint: disable-msg=W0613
+        raise ValueError("%s is a YAML file, please install PyYAML" % file_attr)
+
+    # If file does not exist, provide an error function instead
     if os.path.exists(data_file_path) is False:
         test_name = mk_test_name(name, "error")
-        add_test(cls, test_name, _raise_ve, None)
+        add_test(cls, test_name, _raise_not_exists_error, None)
+        return
+
+    # Load the data as YAML or JSON
+    if data_file_path.endswith((".yml", ".yaml")):
+        # Don't have YAML but want to use YAML file.
+        if not _have_yaml:
+            test_name = mk_test_name(name, "error")
+            add_test(cls, test_name, _raise_need_yaml_error, None)
+            return
+        data = yaml.load(open(data_file_path).read())
     else:
         data = json.loads(open(data_file_path).read())
-        for i, elem in enumerate(data):
-            if isinstance(data, dict):
-                key, value = elem, data[elem]
-                test_name = mk_test_name(name, key, i)
-            elif isinstance(data, list):
-                value = elem
-                test_name = mk_test_name(name, value, i)
-            if isinstance(value, dict):
-                add_test(cls, test_name, func, **value)
-            else:
-                add_test(cls, test_name, func, value)
+
+    # Add tests from data
+    for i, elem in enumerate(data):
+        if isinstance(data, dict):
+            key, value = elem, data[elem]
+            test_name = mk_test_name(name, key, i)
+        elif isinstance(data, list):
+            value = elem
+            test_name = mk_test_name(name, value, i)
+        if isinstance(value, dict):
+            add_test(cls, test_name, func, **value)
+        else:
+            add_test(cls, test_name, func, value)
 
 
 def ddt(cls):

test/test_data_dict.yaml

@@ -0,0 +1,6 @@
+unsorted_list:
+  - 10
+  - 15
+  - 12
+
+sorted_list: [ 15, 12, 50 ]

test/test_data_dict_dict.yaml

@@ -0,0 +1,19 @@
+positive_integer_range:
+    start: 0
+    end: 2
+    value: 1
+
+negative_integer_range:
+    start: -2
+    end: 0
+    value: -1
+
+positive_real_range:
+    start: 0.0
+    end: 1.0
+    value: 0.5
+
+negative_real_range:
+    start: -1.0
+    end: 0.0
+    value: -0.5

test/test_data_list.yaml

@@ -0,0 +1,2 @@
+- "Hello"
+- "Goodbye"

test/test_example.py

@@ -2,6 +2,17 @@ import unittest
 from ddt import ddt, data, file_data, unpack
 from test.mycode import larger_than_two, has_three_elements, is_a_greeting
 
+try:
+    import yaml
+except Exception:
+    do_not_have_yaml_support = True
+else:
+    do_not_have_yaml_support = False
+    del yaml
+
+# A good-looking decorator
+needs_yaml = unittest.skipIf(do_not_have_yaml_support, "Need YAML to run this test")
+
 
 class Mylist(list):
     pass
@@ -32,17 +43,34 @@ class FooTestCase(unittest.TestCase):
         self.assertGreater(a, b)
 
     @file_data("test_data_dict_dict.json")
-    def test_file_data_dict_dict(self, start, end, value):
+    def test_file_data_json_dict_dict(self, start, end, value):
         self.assertLess(start, end)
         self.assertLess(value, end)
         self.assertGreater(value, start)
 
     @file_data('test_data_dict.json')
-    def test_file_data_dict(self, value):
+    def test_file_data_json_dict(self, value):
         self.assertTrue(has_three_elements(value))
 
     @file_data('test_data_list.json')
-    def test_file_data_list(self, value):
+    def test_file_data_json_list(self, value):
+        self.assertTrue(is_a_greeting(value))
+
+    @needs_yaml
+    @file_data("test_data_dict_dict.yaml")
+    def test_file_data_yaml_dict_dict(self, start, end, value):
+        self.assertLess(start, end)
+        self.assertLess(value, end)
+        self.assertGreater(value, start)
+
+    @needs_yaml
+    @file_data('test_data_dict.yaml')
+    def test_file_data_yaml_dict(self, value):
+        self.assertTrue(has_three_elements(value))
+
+    @needs_yaml
+    @file_data('test_data_list.yaml')
+    def test_file_data_yaml_list(self, value):
         self.assertTrue(is_a_greeting(value))
 
     @data((3, 2), (4, 3), (5, 3))