Merge pull request #43 from aai-institute/main-hotfixes

Make first argument to `@nnbench.parametrize` a single iterable

docs/guides/benchmarks.md

@@ -99,9 +99,9 @@ import nnbench
 from training import prepare_data, train_rf, accuracy
 
 @nnbench.parametrize(
-    {"n_estimators": 10, "max_depth": 2},
+    ({"n_estimators": 10, "max_depth": 2},
     {"n_estimators": 50, "max_depth": 5},
-    {"n_estimators": 100, "max_depth": 10}
+    {"n_estimators": 100, "max_depth": 10})
 )
 def benchmark_accuracy(n_estimators: int, max_depth: int, random_state: int) -> float:
     X_train, X_test, y_train, y_test = prepare_data()

src/nnbench/core.py

@@ -102,11 +102,11 @@ def decorator(fun: Callable) -> Benchmark:
 
 
 def parametrize(
-    *parameters: dict[str, Any],
+    parameters: Iterable[dict[str, Any]],
     setUp: Callable[..., None] = NoOp,
     tearDown: Callable[..., None] = NoOp,
     tags: tuple[str, ...] = (),
-) -> list[Benchmark] | Callable[[Callable], list[Benchmark]]:
+) -> Callable[[Callable], list[Benchmark]]:
     """
     Define a family of benchmarks over a function with varying parameters.
 
@@ -116,7 +116,7 @@ def parametrize(
 
     Parameters
     ----------
-    *parameters: dict[str, Any]
+    parameters: Iterable[dict[str, Any]]
         The different sets of parameters defining the benchmark family.
     setUp: Callable[..., None]
         A setup hook to run before each of the benchmarks.
@@ -127,9 +127,8 @@ def parametrize(
 
     Returns
     -------
-    list[Benchmark] | Callable[[Callable], list[Benchmark]]
-        The resulting benchmark family (if no arguments were given), or a parametrized decorator
-        returning the benchmark family.
+    Callable[[Callable], list[Benchmark]]
+        A parametrized decorator returning the benchmark family.
     """
 
     def decorator(fn: Callable) -> list[Benchmark]:
@@ -150,7 +149,7 @@ def product(
     tearDown: Callable[..., None] = NoOp,
     tags: tuple[str, ...] = (),
     **iterables: Iterable,
-) -> list[Benchmark] | Callable[[Callable], list[Benchmark]]:
+) -> Callable[[Callable], list[Benchmark]]:
     """
     Define a family of benchmarks over a cartesian product of one or more iterables.
 
@@ -171,9 +170,8 @@ def product(
 
     Returns
     -------
-    list[Benchmark] | Callable[[Callable], list[Benchmark]]
-        The resulting benchmark family (if no arguments were given), or a parametrized decorator
-        returning the benchmark family.
+    Callable[[Callable], list[Benchmark]]
+        A parametrized decorator returning the benchmark family.
     """
 
     def decorator(fn: Callable) -> list[Benchmark]: