Make index_len argument to idata optional

ddt.py

@@ -90,20 +90,30 @@ def data(*values):
     Should be added to methods of instances of ``unittest.TestCase``.
 
     """
-    return idata(values, len(str(len(values))))
+    return idata(values)
 
 
-def idata(iterable, index_len):
+def idata(iterable, index_len=None):
     """
     Method decorator to add to your test methods.
 
     Should be added to methods of instances of ``unittest.TestCase``.
 
+    :param iterable: iterable of the values to provide to the test function.
+    :param index_len: an optional integer specifying the width to zero-pad the
+        test identifier indices to.  If not provided, this will add the fewest
+        zeros necessary to make all identifiers the same length.
     """
+    if index_len is None:
+        # Avoid consuming a one-time-use generator.
+        iterable = tuple(iterable)
+        index_len = len(str(len(iterable)))
+
     def wrapper(func):
         setattr(func, DATA_ATTR, iterable)
         setattr(func, INDEX_LEN, index_len)
         return func
+
     return wrapper
 
 
@@ -371,4 +381,4 @@ def wrapper(cls):
 
     # ``arg`` is the unittest's test class when decorating with ``@ddt`` while
     # it is ``None`` when decorating a test class with ``@ddt(k=v)``.
-    return wrapper(arg) if inspect.isclass(arg) else wrapper
+    return wrapper(arg) if inspect.isclass(arg) else wrapper

test/test_example.py

@@ -1,6 +1,7 @@
+import itertools
 import unittest
 
-from ddt import ddt, data, file_data, unpack
+from ddt import ddt, data, file_data, idata, unpack
 from test.mycode import larger_than_two, has_three_elements, is_a_greeting
 
 try:
@@ -64,6 +65,12 @@ def test_greater(self, value):
         a, b = value
         self.assertGreater(a, b)
 
+    @idata(itertools.product([0, 1, 2], [3, 4, 5]))
+    def test_iterable_argument(self, value):
+        first_value, second_value = value
+        self.assertLessEqual(first_value, 2)
+        self.assertGreaterEqual(second_value, 3)
+
     @data(annotated2([2, 1], 'Test_case_1', """Test docstring 1"""),
           annotated2([10, 5], 'Test_case_2', """Test docstring 2"""))
     def test_greater_with_name_docstring(self, value):

test/test_functional.py

@@ -9,7 +9,7 @@
 except ImportError:
     import mock
 
-from ddt import ddt, data, file_data, TestNameFormat
+from ddt import ddt, data, file_data, idata, TestNameFormat
 
 from test.mycode import has_three_elements
 
@@ -185,6 +185,97 @@ def test_ddt_format_test_name_default():
         assert ("test_something_{}_{}".format(i, d) in tests)
 
 
+def test_idata_single_argument():
+    """Test that the single-argument form of ``idata`` works."""
+    payload = [5, 12, 13]
+
+    @ddt
+    class Dummy(object):
+        """Dummy class to test that the ``idata(iterable)`` decorator works."""
+        @idata(payload)
+        def test_something(self, value):
+            return value
+
+    tests = list(filter(_is_test, Dummy.__dict__))
+    assert len(tests) == len(payload)
+
+    expected_tests = [
+        "test_something_{:1d}_{}".format(i + 1, v) for i, v in enumerate(payload)
+    ]
+    assert sorted(tests) == sorted(expected_tests)
+
+
+def test_idata_automatic_zero_padding():
+    """
+    Test that the single-argument form of ``idata`` zero-pads its keys so the
+    lengths all match
+    """
+    payload = range(15)
+
+    @ddt
+    class Dummy(object):
+        """Dummy class to test that the ``idata(iterable)`` decorator works."""
+        @idata(payload)
+        def test_something(self, value):
+            return value
+
+    tests = list(filter(_is_test, Dummy.__dict__))
+    assert len(tests) == len(payload)
+
+    expected_tests = [
+        "test_something_{:02d}_{}".format(i + 1, v) for i, v in enumerate(payload)
+    ]
+    assert sorted(tests) == sorted(expected_tests)
+
+
+def test_idata_override_index_len():
+    """
+    Test that overriding ``index_len`` in ``idata`` can allow additional
+    zero-padding to be added.
+    """
+    payload = [4, 2, 1]
+
+    @ddt
+    class Dummy(object):
+        @idata(payload, index_len=2)
+        def test_something(self, value):
+            return value
+
+    tests = list(filter(_is_test, Dummy.__dict__))
+    assert len(tests) == len(payload)
+
+    expected_tests = [
+        "test_something_{:02d}_{}".format(i + 1, v) for i, v in enumerate(payload)
+    ]
+    assert sorted(tests) == sorted(expected_tests)
+
+
+def test_idata_consumable_iterator():
+    """
+    Test that using ``idata`` with a consumable iterator still generates the
+    expected tests.
+    """
+    payload = [51, 78, 2]
+
+    def consumable_iterator():
+        # Not using `yield from` for Python 2.7.
+        for i in payload:
+            yield i
+
+    @ddt
+    class Dummy(object):
+        @idata(consumable_iterator())
+        def test_something(self, value):
+            return value
+
+    tests = list(filter(_is_test, Dummy.__dict__))
+
+    expected_tests = [
+        "test_something_{:1d}_{}".format(i + 1, v) for i, v in enumerate(payload)
+    ]
+    assert sorted(tests) == sorted(expected_tests)
+
+
 def test_file_data_test_creation():
     """
     Test that the ``file_data`` decorator creates two tests