From f24d4e49638868df9f04e19c7982270e5e6daabe Mon Sep 17 00:00:00 2001
From: Daniel Hahler <git@thequod.de>
Date: Tue, 12 Nov 2019 14:52:19 +0100
Subject: [PATCH] [WIP/RFC] parametrized: ids: support generator/iterator

Fixes https://github.com/pytest-dev/pytest/issues/759
---
 src/_pytest/python.py | 53 ++++++++++++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 16 deletions(-)

diff --git a/src/_pytest/python.py b/src/_pytest/python.py
index c629f770ab9..5c3137c1cce 100644
--- a/src/_pytest/python.py
+++ b/src/_pytest/python.py
@@ -943,13 +943,18 @@ def parametrize(self, argnames, argvalues, indirect=False, ids=None, scope=None)
             function so that it can perform more expensive setups during the
             setup phase of a test rather than at collection time.
 
-        :arg ids: list of string ids, or a callable.
+        :arg ids: sequence of string ids, a generator or a callable.
             If strings, each is corresponding to the argvalues so that they are
             part of the test id. If None is given as id of specific test, the
             automatically generated id for that argument will be used.
-            If callable, it should take one argument (a single argvalue) and return
-            a string or return None. If None, the automatically generated id for that
-            argument will be used.
+
+            If callable, it should take either one argument (a single argvalue)
+            and return a string or return None. If it returns None, the
+            automatically generated id for that argument will be used.
+
+            You can also pass a generator, or iterator like
+            ``itertools.count()``.
+
             If no ids are provided they will be generated automatically from
             the argvalues.
 
@@ -1013,27 +1018,43 @@ def _resolve_arg_ids(self, argnames, ids, parameters, item):
         :rtype: List[str]
         :return: the list of ids for each argname given
         """
-        from _pytest._io.saferepr import saferepr
-
         idfn = None
         if callable(ids):
             idfn = ids
             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})"
-                    fail(
-                        msg.format(func_name, saferepr(id_value), type(id_value)),
-                        pytrace=False,
-                    )
+            ids = self._validate_ids(ids, parameters, func_name)
         ids = idmaker(argnames, parameters, idfn, ids, self.config, item=item)
         return ids
 
+    def _validate_ids(self, ids, parameters, func_name):
+        try:
+            len_ids = len(ids)
+        except TypeError:
+            try:
+                it = iter(ids)
+            except TypeError:
+                raise TypeError("ids must be a callable, sequence or generator")
+            else:
+                import itertools
+
+                ids = list(itertools.islice(it, len(parameters)))
+
+        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):
+                from _pytest._io.saferepr import saferepr
+
+                msg = "In {}: ids must be list of strings, found: {} (type: {!r})"
+                fail(
+                    msg.format(func_name, saferepr(id_value), type(id_value)),
+                    pytrace=False,
+                )
+        return ids
+
     def _resolve_arg_value_types(self, argnames, indirect):
         """Resolves if each parametrized argument must be considered a parameter to a fixture or a "funcarg"
         to the function, based on the ``indirect`` parameter of the parametrized() call.