Repository: aai-institute/nnbench

Issue:
Improve ensure unique keys in `Contexts` upon merge.
Follow up from #92

- [x] Define and add a behaviour for non-unique keys in `Context.add` and `Context.update`
- [x] Remove the checks for duplicate keys from the `runner.run` method.

Discussion:

Patch:
diff --git a/src/nnbench/context.py b/src/nnbench/context.py
index b0ad493..4ea8f27 100644
--- a/src/nnbench/context.py
+++ b/src/nnbench/context.py
@@ -88,7 +88,10 @@ class GitEnvironmentInfo:
                 git = "git"
 
             return subprocess.run(  # nosec: B603
-                [git, *args], stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="utf-8"
+                [git, *args],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                encoding="utf-8",
             )
 
         result: dict[str, Any] = {
@@ -246,7 +249,7 @@ class Context:
         str
             Iterator over the context dictionary keys.
         """
-        for k, v in self._ctx_items(d=self._data, prefix="", sep=sep):
+        for k, _ in self._ctx_items(d=self._data, prefix="", sep=sep):
             yield k
 
     def values(self) -> Iterator[Any]:
@@ -258,7 +261,7 @@ class Context:
         Any
             Iterator over all values in the context dictionary.
         """
-        for k, v in self._ctx_items(d=self._data, prefix="", sep=""):
+        for _, v in self._ctx_items(d=self._data, prefix="", sep=""):
             yield v
 
     def items(self, sep: str = ".") -> Iterator[tuple[str, Any]]:
@@ -277,7 +280,7 @@ class Context:
         """
         yield from self._ctx_items(d=self._data, prefix="", sep=sep)
 
-    def add(self, provider: ContextProvider) -> None:
+    def add(self, provider: ContextProvider, replace: bool = False) -> None:
         """
         Adds data from a provider to the context.
 
@@ -285,10 +288,12 @@ class Context:
         ----------
         provider : ContextProvider
             The provider to inject into this context.
+        replace : bool
+            Whether to replace existing context values upon key collision. Raises ValueError otherwise.
         """
-        self._data.update(provider())
+        self.update(Context.make(provider()), replace=replace)
 
-    def update(self, other: "Context") -> None:
+    def update(self, other: "Context", replace: bool = False) -> None:
         """
         Updates the context.
 
@@ -296,7 +301,18 @@ class Context:
         ----------
         other : Context
             The other context to update this context with.
+        replace : bool
+            Whether to replace existing context values upon key collision. Raises ValueError otherwise.
+
+        Raises
+        ------
+        ValueError
+            If ``other contains top-level keys already present in the context and ``replace=False``.
         """
+        duplicates = set(self.keys()) & set(other.keys())
+        if not replace and duplicates:
+            dupe, *_ = duplicates
+            raise ValueError(f"got multiple values for context key {dupe!r}")
         self._data.update(other._data)
 
     @staticmethod
diff --git a/src/nnbench/runner.py b/src/nnbench/runner.py
index d3417b5..56d98d3 100644
--- a/src/nnbench/runner.py
+++ b/src/nnbench/runner.py
@@ -215,11 +215,6 @@ class BenchmarkRunner:
             A JSON output representing the benchmark results. Has two top-level keys, "context"
             holding the context information, and "benchmarks", holding an array with the
             benchmark results.
-
-        Raises
-        ------
-        ValueError
-            If any context value is provided more than once.
         """
         if not self.benchmarks:
             self.collect(path_or_module, tags)
@@ -241,16 +236,8 @@ class BenchmarkRunner:
             ctx = context
         else:
             ctx = Context()
-
             for provider in context:
-                ctx_cand = Context.make(provider())
-
-                # we do not allow multiple values for a context key.
-                duplicates = set(ctx.keys()) & set(ctx_cand.keys())
-                if duplicates:
-                    dupe, *_ = duplicates
-                    raise ValueError(f"got multiple values for context key {dupe!r}")
-                ctx.update(ctx_cand)
+                ctx.add(provider)
 
         results: list[dict[str, Any]] = []
         for benchmark in self.benchmarks:


Test patch:
diff --git a/tests/test_context.py b/tests/test_context.py
index 7c73d6d..f6847fe 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -1,3 +1,5 @@
+import pytest
+
 from nnbench.context import Context, CPUInfo, GitEnvironmentInfo, PythonInfo
 
 
@@ -71,17 +73,30 @@ def test_context_items():
     assert set(ctx.items()) == expected_items
 
 
-def test_update_with_context():
+def test_context_update_with_key_collision():
+    ctx = Context({"a": 1, "b": 2})
+    with pytest.raises(ValueError, match=r".*multiple values for context.*"):
+        ctx.update(Context.make({"a": 3, "c": 4}))
+
+
+def test_context_update_duplicate_with_replace():
     ctx = Context({"a": 1, "b": 2})
-    ctx.update(Context.make({"a": 3, "c": 4}))
+    ctx.update(Context.make({"a": 3, "c": 4}), replace=True)
     expected_dict = {"a": 3, "b": 2, "c": 4}
     assert ctx._data == expected_dict
 
 
+def test_update_with_context():
+    ctx = Context({"a": 1, "b": 2})
+    ctx.update(Context.make({"c": 4}))
+    expected_dict = {"a": 1, "b": 2, "c": 4}
+    assert ctx._data == expected_dict
+
+
 def test_add_with_provider():
     ctx = Context({"a": 1, "b": 2})
-    ctx.add(lambda: {"a": 3, "c": 4})
-    expected_dict = {"a": 3, "b": 2, "c": 4}
+    ctx.add(lambda: {"c": 4})
+    expected_dict = {"a": 1, "b": 2, "c": 4}
     assert ctx._data == expected_dict