Add methods to insert and append (in)to the TraceParameterDefinitionMap

Keysight · Nov 4, 2022 · b3408f2 · b3408f2
1 parent f6f941d
commit b3408f2
Show file tree

Hide file tree

Showing 4 changed files with 121 additions and 14 deletions.
diff --git a/tests/test_header.py b/tests/test_header.py
@@ -90,6 +90,10 @@ def test_trace_param_defs_append_errors(self):
             trace_parameter_definitions.popitem()
         with self.assertRaises(TypeError):
             trace_parameter_definitions.clear()
+        with self.assertRaises(TypeError):
+            trace_parameter_definitions.append('input', ParameterType.BYTE, 16)
+        with self.assertRaises(TypeError):
+            trace_parameter_definitions.insert('output', ParameterType.BYTE, 16, 0)
 
         # Shallow copies still share references to the same trace set parameters as the original,
         # and should therefore not be modifiable if the original isn't

diff --git a/tests/test_parametermap.py b/tests/test_parametermap.py
@@ -109,6 +109,14 @@ def create_parameterdefinitionmap() -> TraceParameterDefinitionMap:
         param_map['中文'] = TraceParameterDefinition(ParameterType.STRING, 15, 29)
         return param_map
 
+    @staticmethod
+    def create_std_parameterdefinitionmap() -> TraceParameterDefinitionMap:
+        param_map = TraceParameterDefinitionMap()
+        param_map['INPUT'] = TraceParameterDefinition(ParameterType.BYTE, 16, 0)
+        param_map['OUTPUT'] = TraceParameterDefinition(ParameterType.BYTE, 16, 16)
+        param_map['KEY'] = TraceParameterDefinition(ParameterType.BYTE, 16, 32)
+        return param_map
+
     @staticmethod
     def create_traceparametermap() -> TraceParameterMap:
         param_map = TraceParameterMap()
@@ -134,6 +142,34 @@ def test_from_trace_params(self):
         map_from_trace_params = TraceParameterDefinitionMap.from_trace_parameter_map(param_map)
         self.assertDictEqual(self.create_parameterdefinitionmap(), map_from_trace_params)
 
+    def test_append(self):
+        map_from_append = TraceParameterDefinitionMap()
+        map_from_append.append('IN', ParameterType.BYTE, 16)
+        map_from_append.append('TITLE', ParameterType.STRING, 13)
+        map_from_append.append('中文', ParameterType.STRING, 15)
+        self.assertDictEqual(self.create_parameterdefinitionmap(), map_from_append)
+
+        map_from_std_append = TraceParameterDefinitionMap()
+        map_from_std_append.append_std('INPUT', 16)
+        map_from_std_append.append_std('OUTPUT', 16)
+        map_from_std_append.append_std('KEY', 16)
+        self.assertDictEqual(self.create_std_parameterdefinitionmap(), map_from_std_append)
+
+    def test_insert(self):
+        map_from_insert = TraceParameterDefinitionMap()
+        map_from_insert.insert('TITLE', ParameterType.STRING, 13, 0)
+        with self.assertWarns(UserWarning):
+            map_from_insert.insert('中文', ParameterType.STRING, 15, 10)
+        map_from_insert.insert('IN', ParameterType.BYTE, 16, 0)
+        self.assertDictEqual(self.create_parameterdefinitionmap(), map_from_insert)
+
+        map_from_std_insert = TraceParameterDefinitionMap()
+        map_from_std_insert.insert_std('INPUT', 16, 0)
+        with self.assertWarns(UserWarning):
+            map_from_std_insert.insert_std('KEY', 16, 9)
+        map_from_std_insert.insert_std('OUTPUT', 16, 16)
+        self.assertDictEqual(self.create_std_parameterdefinitionmap(), map_from_std_insert)
+
 
 class TestTraceParameterMap(TestCase):
     CAFEBABE = bytes.fromhex('cafebabedeadbeef0102030405060708')

diff --git a/trsfile/engine/trs.py b/trsfile/engine/trs.py
@@ -173,21 +173,18 @@ def update_headers_with_traces_metadata(self, traces: List[Trace]) -> None:
             if len(set([len(trace.parameters.serialize()) for trace in traces])) > 1:
                 raise TypeError('Traces have different data length, this is not supported in TRS files')
 
-            data_length = len(traces[0].parameters.serialize())
+            headers_updates[Header.LENGTH_DATA] = len(traces[0].parameters.serialize())
 
-            # Add a TraceParameterDefinitionMap if none is present, and verify its validity if one is present
-            if Header.TRACE_PARAMETER_DEFINITIONS not in self.headers:
-                headers_updates[Header.TRACE_PARAMETER_DEFINITIONS] = \
-                    TraceParameterDefinitionMap.from_trace_parameter_map(traces[0].parameters)
-
-            headers_updates[Header.LENGTH_DATA] = data_length
-
-        # Verify that each trace confirms to the traceset's TraceParameterDefinitionMap
-        for index, trace in enumerate(traces):
-            if not trace.parameters.matches(self.headers[Header.TRACE_PARAMETER_DEFINITIONS]):
-                raise TypeError(f"The parameters of trace #{index} do not match the trace set's definitions.\n"
-                                f"Please make sure the trace parameters match those of the other traces in type, "
-                                f"size and name.")
+        # Add a TraceParameterDefinitionMap if none is present, and verify its validity if one is present
+        if Header.TRACE_PARAMETER_DEFINITIONS not in self.headers:
+            headers_updates[Header.TRACE_PARAMETER_DEFINITIONS] = \
+                TraceParameterDefinitionMap.from_trace_parameter_map(traces[0].parameters)
+        else:
+            for index, trace in enumerate(traces):
+                if not trace.parameters.matches(self.headers[Header.TRACE_PARAMETER_DEFINITIONS]):
+                    raise TypeError(f"The parameters of trace #{index} do not match the trace set's definitions.\n"
+                                    f"Please make sure the trace parameters match those of the other traces in type, "
+                                    f"size and name.")
 
         if self.headers[Header.SAMPLE_CODING] is None:
             if len(set([trace.sample_coding for trace in traces])) > 1:

diff --git a/trsfile/parametermap.py b/trsfile/parametermap.py
@@ -293,6 +293,76 @@ def __setitem__(self, key: str, value: TraceParameterDefinition):
         self._stop_if_locked()
         super().__setitem__(key, value)
 
+    def insert_std(self, name: str, size: int, offset: int):
+        """Insert a trace parameter definition of a StandardTraceParameter into this map in a specified location. If
+        the given offset would put the new TraceParameter in the middle of a parameter already present in the map, the
+        offset is increased to put add the new parameter after that existing one instead. Any parameters already present
+        in the map that have a greater or equal offset than the new parameter, have their offset increased to make space
+        for the new parameter.
+
+        :param name: The name of the TraceParameter for which to add a definition. This name must match that of a
+                    StandardTraceParameter
+        :param size: The size of the TraceParameter, in number of values of its type
+        :param offset: The offset of the TraceParameter, in bytes
+        """
+        try:
+            type = StandardTraceParameters.from_identifier(name).parameter_type
+            self.insert(name, type, size, offset)
+        except ValueError:
+            raise ValueError(f"No StandardTraceParameter found with name '{name}'. Either specify a type in this "
+                             f"insert or correct the name to match a standard trace parameter.")
+
+    def insert(self, name: str, type: ParameterType, size: int, offset: int):
+        """Insert a trace parameter definition into this map in a specified location. If the given offset would put the
+        new TraceParameter in the middle of a parameter already present in the map, the offset is increased to put add
+        the new parameter after that existing one instead. Any parameters already present in the map that have a greater
+        or equal offset than the new parameter, have their offset increased to make space for the new parameter.
+
+        :param name: The name of the TraceParameter for which to add a definition
+        :param type: The type of the TraceParameter for which to add a definition
+        :param size: The size of the TraceParameter, in number of values of its type
+        :param offset: The offset of the TraceParameter, in bytes
+        """
+        self._stop_if_locked()
+        params_to_move_back = []
+        for key, param in self.items():
+            if param.offset >= offset:
+                param.offset += size * type.byte_size
+                params_to_move_back.append(key)
+            elif param.offset + param.length * param.param_type.byte_size > offset:
+                offset = param.offset + param.length * param.param_type.byte_size
+                warnings.warn("Given offset would put a parameter inside another trace parameter.\n"
+                              f"Increased the offset of the inserted parameter definition to {offset} to prevent this.")
+
+        new_definition = TraceParameterDefinition(type, size, offset)
+        self.__setitem__(name, new_definition)
+        for param in params_to_move_back:
+            self.move_to_end(param)
+
+    def append_std(self, name: str, size: int):
+        """Append a trace parameter definition of a StandardTraceParameter to this map. The parameter wil be added after
+       all parameter definitions already in the map.
+
+       :param name: The name of the TraceParameter for which to add a definition. This name must match that of a
+                    StandardTraceParameter
+       :param size: The size of the TraceParameter, in number of values of its type"""
+        try:
+            type = StandardTraceParameters.from_identifier(name).parameter_type
+            self.append(name, type, size)
+        except ValueError:
+            raise ValueError(f"No StandardTraceParameter found with name '{name}'. Either specify a type in this "
+                             f"append or correct the name to match a standard trace parameter.")
+
+    def append(self, name: str, type: ParameterType, size: int):
+        """Append a trace parameter definition to this map. The parameter wil be added after all parameter definitions
+        already in the map.
+
+        :param name: The name of the TraceParameter for which to add a definition
+        :param type: The type of the TraceParameter for which to add a definition
+        :param size: The size of the TraceParameter, in number of values of its type"""
+        new_definition = TraceParameterDefinition(type, size, self.get_total_size())
+        self.__setitem__(name, new_definition)
+
     @staticmethod
     def deserialize(raw: BytesIO) -> TraceParameterDefinitionMap:
         result = TraceParameterDefinitionMap()