From 87a81837bc76db1f1eeac51214b2b13ae2f942e2 Mon Sep 17 00:00:00 2001
From: Josse Van Delm <josse.vandelm@gmail.com>
Date: Fri, 24 Jan 2025 17:16:18 +0100
Subject: [PATCH] Add typing

---
 util/tracing/merge_json.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/util/tracing/merge_json.py b/util/tracing/merge_json.py
index a1558beb..86915606 100644
--- a/util/tracing/merge_json.py
+++ b/util/tracing/merge_json.py
@@ -7,7 +7,9 @@
 
 def merge_json(file_list: Sequence[str], output_file: str):
     prefix = ""
-    traces = defaultdict(lambda: defaultdict(list))
+    traces: defaultdict[
+        int, defaultdict[int, list[dict[str, int | float | str | None]]]
+    ] = defaultdict(lambda: defaultdict(list))
     for filename in file_list:
         match = re.match(r"(.*)_trace_chip_(\d{2})_hart_(\d{5})_perf\.json", filename)
         if match:
@@ -29,7 +31,7 @@ def merge_json(file_list: Sequence[str], output_file: str):
     # Convert the defaultdict to a nested list structure
     result = [
         [chip_harts[hart_id] for hart_id in sorted(chip_harts)]
-        for chip_id, chip_harts in sorted(traces.items())
+        for _, chip_harts in sorted(traces.items())
     ]
 
     # Use prefix for output file if not explicitly provided
@@ -48,18 +50,17 @@ def merge_json(file_list: Sequence[str], output_file: str):
         "files",
         metavar="FILE",
         nargs="+",
-        help=(
-            "List of JSON files to merge.",
-            "Filenames must match the format",
-            "'<prefix>_trace_chip_<chip_id>_hart_<hart_id>_perf.json'.",
-        ),
+        help="List of JSON files to merge."
+        "Filenames must match the format"
+        "'<prefix>_trace_chip_<chip_id>_hart_<hart_id>_perf.json'.",
     )
     parser.add_argument(
         "-o",
         "--output",
         metavar="OUTPUT_FILE",
         default=None,
-        help="Name of the output JSON file. Defaults to '<prefix>_aggregated.json' based on input filenames.",
+        help="Name of the output JSON file."
+        "Defaults to '<prefix>_aggregated.json' based on input filenames.",
     )
 
     args = parser.parse_args()