pw_trace: Add mapping format to pw_trace decoder

Adding new format logic + documentation

Bug: 269176238
Test: Added unit test for one and two mapped fields
Change-Id: I53d38dc664c12be86997a1884b9ace594795dd9c
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/128771
Reviewed-by: Victor Berchet <berchet@google.com>
Commit-Queue: Aditya Chitnis <chitnis@google.com>
Reviewed-by: Keir Mierle <keir@google.com>

3 files changed, 165 insertions, 0 deletions

diff --git a/pw_trace/docs.rst b/pw_trace/docs.rst
index 0791f7ffa..7e328aa66 100644
--- a/pw_trace/docs.rst
+++ b/pw_trace/docs.rst
@@ -199,6 +199,15 @@ Currently the included python tool supports a few different options for
   can be used to either provide a single value type, or provide multiple
   different values with a variety of types. Options for format string types can
   be found here: https://docs.python.org/3/library/struct.html#format-characters
+- *@pw_py_map_fmt:* - Interprets the string after the ":" as a dictionary
+  relating the data field name to the python struct format string. Once
+  collected, the format strings are concatenated and used to unpack the data
+  elements as above::
+
+    // Example
+    data_format_string = "@pw_py_map_fmt:{Field: l, Field2: l }"
+    data = 0x14000000000000001400000000000000 (little endian)
+    args = {Field: 20, Field2:20}
 
 .. tip::
 
diff --git a/pw_trace/py/pw_trace/trace.py b/pw_trace/py/pw_trace/trace.py
index 4a5b91eda..dbd2b7d16 100755
--- a/pw_trace/py/pw_trace/trace.py
+++ b/pw_trace/py/pw_trace/trace.py
@@ -69,6 +69,34 @@ def event_has_trace_id(event_type):
     }
 
 
+def decode_map_fmt_args(event):
+    """Decodes the trace's event data for map-formatted data"""
+    args = {}
+    fmt_list = event.data_fmt[len("@pw_py_map_fmt:") :].strip("{}").split(",")
+    fmt_bytes = ''
+    fields = []
+    try:
+        for pair in fmt_list:
+            (field, value) = (s.strip() for s in pair.split(":"))
+            fields.append(field)
+            fmt_bytes += value
+    except ValueError:
+        args["error"] = f"Invalid map format {event.data_fmt}"
+    else:
+        try:
+            # needed in case the buffer is larger than expected
+            assert struct.calcsize(fmt_bytes) == len(event.data)
+            items = struct.unpack_from(fmt_bytes, event.data)
+            for i, item in enumerate(items):
+                args[fields[i]] = item
+        except (AssertionError, struct.error):
+            args["error"] = (
+                f"Mismatched struct/data format {event.data_fmt}"
+                f" data {event.data.hex()}"
+            )
+    return args
+
+
 def generate_trace_json(events: Iterable[TraceEvent]):
     """Generates a list of JSON lines from provided trace events."""
     json_lines = []
@@ -150,6 +178,8 @@ def generate_trace_json(events: Iterable[TraceEvent]):
                 for i, item in enumerate(items):
                     args["data_" + str(i)] = item
                 line["args"] = args
+            elif event.data_fmt.startswith("@pw_py_map_fmt:"):
+                line["args"] = decode_map_fmt_args(event)
             else:
                 line["args"] = {"data": event.data.hex()}
 
diff --git a/pw_trace/py/trace_test.py b/pw_trace/py/trace_test.py
index e3196d75b..9cec920bb 100755
--- a/pw_trace/py/trace_test.py
+++ b/pw_trace/py/trace_test.py
@@ -211,6 +211,132 @@ class TestTraceGenerateJson(unittest.TestCase):
             },
         )
 
+    def test_generate_json_data_map_fmt_single(self):
+        event = trace.TraceEvent(
+            event_type=trace.TraceType.INSTANTANEOUS,
+            module="module",
+            label="label",
+            timestamp_us=10,
+            has_data=True,
+            data_fmt="@pw_py_map_fmt:{Field:l}",
+            data=struct.pack("l", 20),
+        )
+        json_lines = trace.generate_trace_json([event])
+        self.assertEqual(1, len(json_lines))
+        self.assertEqual(
+            json.loads(json_lines[0]),
+            {
+                "ph": "I",
+                "pid": "module",
+                "name": "label",
+                "ts": 10,
+                "s": "p",
+                "args": {"Field": 20},
+            },
+        )
+
+    def test_generate_json_data_map_fmt_multi(self):
+        event = trace.TraceEvent(
+            event_type=trace.TraceType.INSTANTANEOUS,
+            module="module",
+            label="label",
+            timestamp_us=10,
+            has_data=True,
+            data_fmt="@pw_py_map_fmt:{Field: l, Field2: l }",
+            data=struct.pack("ll", 20, 40),
+        )
+        json_lines = trace.generate_trace_json([event])
+        self.assertEqual(1, len(json_lines))
+        self.assertEqual(
+            json.loads(json_lines[0]),
+            {
+                "ph": "I",
+                "pid": "module",
+                "name": "label",
+                "ts": 10,
+                "s": "p",
+                "args": {"Field": 20, "Field2": 40},
+            },
+        )
+
+    def test_generate_error_json_data_map_bad_fmt(self):
+        event = trace.TraceEvent(
+            event_type=trace.TraceType.INSTANTANEOUS,
+            module="module",
+            label="label",
+            timestamp_us=10,
+            has_data=True,
+            data_fmt="@pw_py_map_fmt:{Field;l,Field2;l}",
+            data=struct.pack("ll", 20, 40),
+        )
+        json_lines = trace.generate_trace_json([event])
+        self.assertEqual(1, len(json_lines))
+        self.assertEqual(
+            json.loads(json_lines[0]),
+            {
+                "ph": "I",
+                "pid": "module",
+                "name": "label",
+                "ts": 10,
+                "s": "p",
+                "args": {"error": f"Invalid map format {event.data_fmt}"},
+            },
+        )
+
+    def test_generate_error_json_data_map_invalid_small_buffer(self):
+        event = trace.TraceEvent(
+            event_type=trace.TraceType.INSTANTANEOUS,
+            module="module",
+            label="label",
+            timestamp_us=10,
+            has_data=True,
+            data_fmt="@pw_py_map_fmt:{Field:l,Field2:l}",
+            data=struct.pack("l", 20),
+        )
+        json_lines = trace.generate_trace_json([event])
+        self.assertEqual(1, len(json_lines))
+        self.assertEqual(
+            json.loads(json_lines[0]),
+            {
+                "ph": "I",
+                "pid": "module",
+                "name": "label",
+                "ts": 10,
+                "s": "p",
+                "args": {
+                    "error": f"Mismatched struct/data format {event.data_fmt} "
+                    f"data {event.data.hex()}"
+                },
+            },
+        )
+
+    def test_generate_error_json_data_map_invalid_large_buffer(self):
+        event = trace.TraceEvent(
+            event_type=trace.TraceType.INSTANTANEOUS,
+            module="module",
+            label="label",
+            timestamp_us=10,
+            has_data=True,
+            data_fmt="@pw_py_map_fmt:{Field:H,Field2:H}",
+            data=struct.pack("ll", 20, 40),
+        )
+        json_lines = trace.generate_trace_json([event])
+        self.assertEqual(1, len(json_lines))
+        self.assertEqual(
+            json.loads(json_lines[0]),
+            {
+                "ph": "I",
+                "pid": "module",
+                "name": "label",
+                "ts": 10,
+                "s": "p",
+                "args": {
+                    "error": f"Mismatched struct/data format {event.data_fmt} "
+                    f"data {event.data.hex()}"
+                },
+            },
+        )
+
 
 if __name__ == '__main__':
     unittest.main()