Add ability to specify render per wire name in trace

Author: mdko

pyrtl/simulation.py

@@ -876,35 +876,46 @@ def tick_segment(self, n, symbol_len, segment_size):
         num_tick = self._tick + str(n)
         return num_tick.ljust(symbol_len * segment_size)
 
-    def render_val(self, w, n, current_val, symbol_len, repr_func):
+    def render_val(self, w, n, current_val, symbol_len, repr_func, repr_per_name):
         if w is not self.prev_wire:
             self.prev_wire = w
             self.prior_val = current_val
-        out = self._render_val_with_prev(w, n, current_val, symbol_len, repr_func)
+        out = self._render_val_with_prev(w, n, current_val, symbol_len, repr_func, repr_per_name)
         self.prior_val = current_val
         return out
 
-    def _render_val_with_prev(self, w, n, current_val, symbol_len, repr_func):
+    def _render_val_with_prev(self, w, n, current_val, symbol_len, repr_func, repr_per_name):
         """Return a string encoding the given value in a waveform.
 
         :param w: The WireVector we are rendering to a waveform
         :param n: An integer from 0 to segment_len-1
         :param current_val: the value to be rendered
         :param symbol_len: and integer for how big to draw the current value
         :param repr_func: function to use for representing the current_val;
-            examples are 'hex', 'oct', 'bin', and 'str' (for decimal). Defaults to 'hex'.
+            examples are 'hex', 'oct', 'bin', 'str' (for decimal), or even the name
+            of an IntEnum class you know the value will belong to. Defaults to 'hex'.
+        :param repr_per_name: Map from signal name to a function that takes in the signal's
+            value and returns a user-defined representation. If a signal name is
+            not found in the map, the argument `repr_func` will be used instead.
 
         Returns a string of printed length symbol_len that will draw the
         representation of current_val.  The input prior_val is used to
         render transitions.
         """
+        def to_str(v):
+            f = repr_per_name.get(w.name)
+            if f is None:
+                return str(repr_func(v))
+            else:
+                return str(f(v))
+
         sl = symbol_len - 1
         if len(w) > 1:
             out = self._revstart
             if current_val != self.prior_val:
-                out += self._x + repr_func(current_val).rstrip('L').ljust(sl)[:sl]
+                out += self._x + to_str(current_val).rstrip('L').ljust(sl)[:sl]
             elif n == 0:
-                out += repr_func(current_val).rstrip('L').ljust(symbol_len)[:symbol_len]
+                out += to_str(current_val).rstrip('L').ljust(symbol_len)[:symbol_len]
             else:
                 out += ' ' * symbol_len
             out += self._revstop
@@ -1132,7 +1143,8 @@ def print_trace_strs(time):
 
     def render_trace(
             self, trace_list=None, file=sys.stdout, render_cls=default_renderer(),
-            symbol_len=5, repr_func=hex, segment_size=5, segment_delim=' ', extra_line=True):
+            symbol_len=5, repr_func=hex, repr_per_name={}, segment_size=5,
+            segment_delim=' ', extra_line=True):
 
         """ Render the trace to a file using unicode and ASCII escape sequences.
 
@@ -1144,6 +1156,9 @@ def render_trace(
             represented value fits.
         :param repr_func: Function to use for representing each value in the trace;
             examples are 'hex', 'oct', 'bin', and 'str' (for decimal). Defaults to 'hex'.
+        :param repr_per_name: Map from signal name to a function that takes in the signal's
+            value and returns a user-defined representation. If a signal name is
+            not found in the map, the argument `repr_func` will be used instead.
         :param segment_size: Traces are broken in the segments of this number of cycles.
         :param segment_delim: The character to be output between segments.
         :param extra_line: A Boolean to determine if we should print a blank line between signals.
@@ -1172,12 +1187,12 @@ def render_trace(
         else:
             self.render_trace_to_text(
                 trace_list=trace_list, file=file, render_cls=render_cls,
-                symbol_len=symbol_len, repr_func=repr_func, segment_size=segment_size,
-                segment_delim=segment_delim, extra_line=extra_line)
+                symbol_len=symbol_len, repr_func=repr_func, repr_per_name=repr_per_name,
+                segment_size=segment_size, segment_delim=segment_delim, extra_line=extra_line)
 
     def render_trace_to_text(
             self, trace_list, file, render_cls,
-            symbol_len, repr_func, segment_size, segment_delim, extra_line):
+            symbol_len, repr_func, repr_per_name, segment_size, segment_delim, extra_line):
 
         renderer = render_cls()
 
@@ -1192,7 +1207,8 @@ def formatted_trace_line(wire, trace):
                     i % segment_size,
                     trace[i],
                     symbol_len,
-                    repr_func)
+                    repr_func,
+                    repr_per_name)
             return heading + trace_line
 
         # default to printing all signals in sorted order
@@ -1212,9 +1228,17 @@ def formatted_trace_line(wire, trace):
                 "if a CompiledSimulation was used.")
 
         if symbol_len is None:
+
+            def to_str(v, name):
+                f = repr_per_name.get(name)
+                if f is None:
+                    return str(repr_func(v))
+                else:
+                    return str(f(v))
+
             maxvallen = 0
-            for trace in self.trace.values():
-                maxvallen = max(maxvallen, max(len(repr_func(v)) for v in trace))
+            for name, trace in self.trace.items():
+                maxvallen = max(maxvallen, max(len(to_str(v, name)) for v in trace))
             symbol_len = maxvallen + 1
 
         # print the 'ruler' which is just a list of 'ticks'

pyrtl/visualization.py

@@ -459,12 +459,18 @@ def block_to_svg(block=None, split_state=True, maintain_arg_order=False):
 #    |__|  |  |\/| |
 #    |  |  |  |  | |___
 
-def trace_to_html(simtrace, trace_list=None, sortkey=None, repr_func=hex):
+def trace_to_html(simtrace, trace_list=None, sortkey=None, repr_func=hex, repr_per_name={}):
     """ Return a HTML block showing the trace.
 
     :param simtrace: A SimulationTrace object
     :param trace_list: (optional) A list of wires to display
     :param sortkey: (optional) The key with which to sort the trace_list
+    :param repr_func: function to use for representing the current_val;
+        examples are 'hex', 'oct', 'bin', 'str' (for decimal), or even the name
+        of an IntEnum class you know the value will belong to. Defaults to 'hex'.
+    :param repr_per_name: Map from signal name to a function that takes in the signal's
+        value and returns a user-defined representation. If a signal name is
+        not found in the map, the argument `repr_func` will be used instead.
     :return: An HTML block showing the trace
     """
 
@@ -506,16 +512,23 @@ def extract(w):
                     wavelist.append(str(value))
                 else:
                     wavelist.append('=')
-                    datalist.append(value)
+                    datalist.append((value, w))
                 last = value
 
+        def to_str(v, name):
+            f = repr_per_name.get(name)
+            if f is None:
+                return str(repr_func(v))
+            else:
+                return str(f(v))
+
         wavestring = ''.join(wavelist)
-        datastring = ', '.join(['"%s"' % repr_func(data) for data in datalist])
+        datastring = ', '.join(['"%s"' % to_str(data, name) for data, name in datalist])
         if len(simtrace._wires[w]) == 1:
             vallens.append(1)  # all are the same length
             return bool_signal_template % (w, wavestring)
         else:
-            vallens.extend([len(repr_func(data)) for data in datalist])
+            vallens.extend([len(to_str(data, name)) for data, name in datalist])
             return int_signal_template % (w, wavestring, datastring)
 
     bool_signal_template = '    { name: "%s",  wave: "%s" },'

tests/test_simulation.py

@@ -180,6 +180,50 @@ def test_decimal_trace(self):
         self.check_rendered_trace(expected, repr_func=str, symbol_len=None)
 
 
+class RenderTraceCustomBase(unittest.TestCase):
+    def setUp(self):
+        pyrtl.reset_working_block()
+
+    def test_custom_repr_per_wire(self):
+        from enum import IntEnum
+
+        class Foo(IntEnum):
+            A = 0
+            B = 1
+            C = 2
+            D = 3
+
+        i = pyrtl.Input(4, 'i')
+        state = pyrtl.Register(max(Foo).bit_length(), name='state')
+        o = pyrtl.Output(name='o')
+        o <<= state
+
+        with pyrtl.conditional_assignment:
+            with i == 0b0001:
+                state.next |= Foo.A
+            with i == 0b0010:
+                state.next |= Foo.B
+            with i == 0b0100:
+                state.next |= Foo.C
+            with i == 0b1000:
+                state.next |= Foo.D
+
+        sim = pyrtl.Simulation()
+        sim.step_multiple({
+            'i': [1, 2, 4, 8, 0]
+        })
+        buff = io.StringIO()
+        sim.tracer.render_trace(file=buff, render_cls=pyrtl.simulation.AsciiWaveRenderer,
+                                extra_line=None, repr_per_name={'state': Foo}, symbol_len=None)
+        expected = (
+            "      -0                            \n"
+            "    i  0x1     x0x2    x0x4    x0x8    x0x0   \n"
+            "    o  0x0             x0x1    x0x2    x0x3   \n"
+            "state  Foo.A           xFoo.B  xFoo.C  xFoo.D \n"
+        )
+        self.assertEqual(buff.getvalue(), expected)
+
+
 class PrintTraceBase(unittest.TestCase):
     # note: doesn't include tests for compact=True because all the tests test that
     def setUp(self):

tests/test_visualization.py

@@ -382,26 +382,50 @@ def test_trace_to_html_repr_func(self):
         )
         self.assertEqual(htmlstring, expected)
 
-    def test_trace_to_html_repr_func_2(self):
-        i = pyrtl.Input(1, 'i')
-        o = pyrtl.Output(2, 'o')
-        o <<= i + 1
+    def test_trace_to_html_repr_per_name(self):
+        from enum import IntEnum
+
+        class Foo(IntEnum):
+            A = 0
+            B = 1
+            C = 2
+            D = 3
+
+        i = pyrtl.Input(4, 'i')
+        state = pyrtl.Register(max(Foo).bit_length(), name='state')
+        o = pyrtl.Output(name='o')
+        o <<= state
+
+        with pyrtl.conditional_assignment:
+            with i == 0b0001:
+                state.next |= Foo.A
+            with i == 0b0010:
+                state.next |= Foo.B
+            with i == 0b0100:
+                state.next |= Foo.C
+            with i == 0b1000:
+                state.next |= Foo.D
 
         sim = pyrtl.Simulation()
-        sim.step_multiple({'i': '0100110'})
-        htmlstring = pyrtl.trace_to_html(sim.tracer, repr_func=bin)
+        sim.step_multiple({
+            'i': [1, 2, 4, 8, 0]
+        })
+
+        htmlstring = pyrtl.trace_to_html(sim.tracer, repr_per_name={'state': Foo})
         expected = (
             '<script type="WaveDrom">\n'
             '{\n'
             '  signal : [\n'
-            '    { name: "i",  wave: "010.1.0" },\n'
-            '    { name: "o",  wave: "===.=.=", data: ["0b1", "0b10", "0b1", "0b10", "0b1"] },\n'
+            '    { name: "i",  wave: "=====", data: ["0x1", "0x2", "0x4", "0x8", "0x0"] },\n'
+            '    { name: "o",  wave: "=.===", data: ["0x0", "0x1", "0x2", "0x3"] },\n'
+            '    { name: "state",  wave: "=.===", data: ["Foo.A", "Foo.B", "Foo.C", "Foo.D"] },\n'
             '  ],\n'
-            '  config: { hscale: 1 }\n'
+            '  config: { hscale: 2 }\n'
             '}\n'
             '</script>\n'
         )
         self.assertEqual(htmlstring, expected)
 
+
 if __name__ == "__main__":
     unittest.main()