parametrize: ids: support iterator

Ref: pytest-dev#759
blueyed · Nov 12, 2019 · b9ddfb3 · b9ddfb3
1 parent 837ac4d
commit b9ddfb3
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 8 deletions.
diff --git a/src/_pytest/python.py b/src/_pytest/python.py
@@ -1032,16 +1032,25 @@ def _resolve_arg_ids(self, argnames, ids, parameters, item, idsetfn):
             ids = None
         if ids:
             func_name = self.function.__name__
-            if len(ids) != len(parameters):
-                msg = "In {}: {} parameter sets specified, with different number of ids: {}"
-                fail(msg.format(func_name, len(parameters), len(ids)), pytrace=False)
-            for id_value in ids:
-                if id_value is not None and not isinstance(id_value, str):
-                    msg = "In {}: ids must be list of strings, found: {} (type: {!r})"
+            try:
+                len_ids = len(ids)
+            except TypeError:
+                len_ids = None
+            if len_ids is not None:
+                if len_ids != len(parameters):
+                    msg = "In {}: {} parameter sets specified, with different number of ids: {}"
                     fail(
-                        msg.format(func_name, saferepr(id_value), type(id_value)),
-                        pytrace=False,
+                        msg.format(func_name, len(parameters), len(ids)), pytrace=False
                     )
+                for id_value in ids:
+                    if id_value is not None and not isinstance(id_value, str):
+                        msg = (
+                            "In {}: ids must be list of strings, found: {} (type: {!r})"
+                        )
+                        fail(
+                            msg.format(func_name, saferepr(id_value), type(id_value)),
+                            pytrace=False,
+                        )
         ids = idmaker(
             argnames, parameters, idfn, ids, self.config, item=item, idsetfn=idsetfn
         )
@@ -1186,6 +1195,8 @@ def _idval(val, argname, idx, idfn, item, config):
 
 
 def _idvalset(idx, parameterset, argnames, idfn, ids, item, config):
+    if hasattr(ids, "__next__"):
+        return next(ids)
     if ids is None or (idx >= len(ids) or ids[idx] is None):
         this_id = [
             _idval(val, argname, idx, idfn, item=item, config=config)

diff --git a/testing/python/metafunc.py b/testing/python/metafunc.py
@@ -1857,3 +1857,23 @@ def test_foo(a):
         )
         result = testdir.runpytest()
         result.assert_outcomes(passed=1)
+
+    def test_parametrize_iterator(self, testdir):
+        testdir.makepyfile(
+            """
+            import itertools
+            import pytest
+
+            @pytest.mark.parametrize("a, b", [(1, 2), (3, 4)], ids=itertools.count(1))
+            def test_foo(a, b):
+                pass
+        """
+        )
+        result = testdir.runpytest("-vv", "-s")
+        result.assert_outcomes(passed=2)
+        result.stdout.fnmatch_lines(
+            [
+                "test_parametrize_iterator.py::test_foo[1] PASSED",
+                "test_parametrize_iterator.py::test_foo[2] PASSED",
+            ]
+        )