Tests: Reduce amount of nesting in benchmarking JSON output

Signed-off-by: Hanno Becker <[email protected]>
pq-code-package · Jan 1, 2025 · f331ee4 · f331ee4
1 parent ceea1f5
commit f331ee4
Showing 1 changed file with 26 additions and 27 deletions.
diff --git a/scripts/lib/mlkem_test.py b/scripts/lib/mlkem_test.py
@@ -277,34 +277,33 @@ def bench(self):
 
         # NOTE: There will only be one items in resultss, as we haven't yet decided how to write both opt/no-opt benchmark results
         for k, results in resultss.items():
-            if results is not None and output is not None and components is False:
-                import json
-
-                with open(output, "w") as f:
-                    v = []
-                    for scheme in results:
-                        schemeStr = str(scheme)
-                        r = results[scheme]
-
-                        # The first 3 lines of the output are expected to be
-                        # keypair cycles=X
-                        # encaps cycles=X
-                        # decaps cycles=X
-
-                        lines = [line for line in r.splitlines() if "=" in line]
-
-                        d = {
-                            k.strip(): int(v) for k, v in (l.split("=") for l in lines)
+            if not (results is not None and output is not None and components is False):
+                continue
+
+            v = []
+            for scheme in results:
+                schemeStr = str(scheme)
+                r = results[scheme]
+
+                # The first 3 lines of the output are expected to be
+                # keypair cycles=X
+                # encaps cycles=X
+                # decaps cycles=X
+
+                lines = [line for line in r.splitlines() if "=" in line]
+
+                d = {k.strip(): int(v) for k, v in (l.split("=") for l in lines)}
+                for primitive in ["keypair", "encaps", "decaps"]:
+                    v.append(
+                        {
+                            "name": f"{schemeStr} {primitive}",
+                            "unit": "cycles",
+                            "value": d[f"{primitive} cycles"],
                         }
-                        for primitive in ["keypair", "encaps", "decaps"]:
-                            v.append(
-                                {
-                                    "name": f"{schemeStr} {primitive}",
-                                    "unit": "cycles",
-                                    "value": d[f"{primitive} cycles"],
-                                }
-                            )
-                    f.write(json.dumps(v))
+                    )
+
+            with open(output, "w") as f:
+                f.write(json.dumps(v))
 
     def all(self):
         func = self.args.func