Support passing generators to parametrize

gundotio · Jan 13, 2023 · 6f9a73d · 6f9a73d
1 parent e4d7d50
commit 6f9a73d
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/__init__.py b/tests/__init__.py
@@ -1,3 +1,5 @@
+from functools import wraps
+
 import pytest
 import pytest_factoryboy
 
@@ -6,9 +8,12 @@
 register = pytest_factoryboy.register
 
 
+@wraps(pytest.mark.kwparametrize)
 def parametrize(*args, **kwargs):
     if not args:
         return pytest.mark.kwparametrize(kwargs)
     if isinstance(args[0], dict):
         return pytest.mark.kwparametrize(*args, **kwargs)
+    if len(args) == 1 and not kwargs:
+        return pytest.mark.kwparametrize(*args[0])
     return pytest.mark.parametrize(*args, **kwargs)