gh-112620: Fix dis error on show_cache with labels (#112621)

iritkatriel · web-flow · commit 162d3d428a83 · 2023-12-03T12:12:49.000Z
diff --git a/Lib/dis.py b/Lib/dis.py
@@ -412,15 +412,14 @@ def _create(cls, op, arg, offset, start_offset, starts_line, line_number,
                co_consts=None, varname_from_oparg=None, names=None,
                labels_map=None, exceptions_map=None):
 
-        label_width = 4 + len(str(len(labels_map)))
         argval, argrepr = cls._get_argval_argrepr(
                                op, arg, offset,
                                co_consts, names, varname_from_oparg, labels_map)
         label = labels_map.get(offset, None)
         instr = Instruction(_all_opname[op], op, arg, argval, argrepr,
                             offset, start_offset, starts_line, line_number,
                             label, positions)
-        instr.label_width = label_width
+        instr.label_width = 4 + len(str(len(labels_map)))
         instr.exc_handler = exceptions_map.get(offset, None)
         return instr
 
@@ -468,12 +467,14 @@ def is_jump_target(self):
         """True if other code jumps to here, otherwise False"""
         return self.label is not None
 
-    def _disassemble(self, lineno_width=3, mark_as_current=False, offset_width=0):
+    def _disassemble(self, lineno_width=3, mark_as_current=False, offset_width=0,
+                           label_width=0):
         """Format instruction details for inclusion in disassembly output.
 
         *lineno_width* sets the width of the line number field (0 omits it)
         *mark_as_current* inserts a '-->' marker arrow as part of the line
         *offset_width* sets the width of the instruction offset field
+        *label_width* sets the width of the label field
         """
         fields = []
         # Column: Source code line number
@@ -488,9 +489,9 @@ def _disassemble(self, lineno_width=3, mark_as_current=False, offset_width=0):
         # Column: Label
         if self.label is not None:
             lbl = f"L{self.label}:"
-            fields.append(f"{lbl:>{self.label_width}}")
+            fields.append(f"{lbl:>{label_width}}")
         else:
-            fields.append(' ' * self.label_width)
+            fields.append(' ' * label_width)
         # Column: Instruction offset from start of code sequence
         if offset_width > 0:
             fields.append(f"{repr(self.offset):>{offset_width}}  ")
@@ -648,6 +649,7 @@ def make_labels_map(original_code, exception_entries):
         return labels_map
 
     labels_map = make_labels_map(original_code, exception_entries)
+    label_width = 4 + len(str(len(labels_map)))
 
     exceptions_map = {}
     for start, end, target, _, _ in exception_entries:
@@ -756,6 +758,7 @@ def _disassemble_bytes(code, lasti=-1, varname_from_oparg=None,
     else:
         offset_width = 0
 
+    label_width = -1
     for instr in _get_instructions_bytes(code, varname_from_oparg, names,
                                          co_consts, linestarts,
                                          line_offset=line_offset,
@@ -774,7 +777,9 @@ def _disassemble_bytes(code, lasti=-1, varname_from_oparg=None,
             # Each CACHE takes 2 bytes
             is_current_instr = instr.offset <= lasti \
                 <= instr.offset + 2 * _get_cache_size(_all_opname[_deoptop(instr.opcode)])
-        print(instr._disassemble(lineno_width, is_current_instr, offset_width),
+        label_width = getattr(instr, 'label_width', label_width)
+        assert label_width >= 0
+        print(instr._disassemble(lineno_width, is_current_instr, offset_width, label_width),
               file=file)
     if exception_entries:
         print("ExceptionTable:", file=file)
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -1896,6 +1896,18 @@ def test_oparg_alias(self):
                                   positions=None)
         self.assertEqual(instruction.arg, instruction.oparg)
 
+    def test_show_caches_with_label(self):
+        def f(x, y, z):
+            if x:
+                res = y
+            else:
+                res = z
+            return res
+
+        output = io.StringIO()
+        dis.dis(f.__code__, file=output, show_caches=True)
+        self.assertIn("L1:", output.getvalue())
+
     def test_baseopname_and_baseopcode(self):
         # Standard instructions
         for name, code in dis.opmap.items():
