gh-144207: Syntax highlighting(theming support) for dis module

Lib/_colorize.py

@@ -200,6 +200,88 @@ class Difflib(ThemeSection):
     reset: str = ANSIColors.RESET
 
 
+@dataclass(frozen=True, kw_only=True)
+class Dis(ThemeSection):
+    label_bg: str = ANSIColors.BACKGROUND_BLUE
+    label_fg: str = ANSIColors.BLACK
+    exception_label: str = ANSIColors.CYAN
+    argument_detail: str = ANSIColors.GREY
+
+    op_stack: str = ANSIColors.BOLD_YELLOW
+    op_load_store: str = ANSIColors.BOLD_CYAN
+    op_call_return: str = ANSIColors.BOLD_MAGENTA
+    op_binary_unary: str = ANSIColors.BOLD_BLUE
+    op_control_flow: str = ANSIColors.BOLD_GREEN
+    op_build: str = ANSIColors.BOLD_WHITE
+    op_exceptions: str = ANSIColors.BOLD_RED
+    op_other: str = ANSIColors.GREY
+
+    reset: str = ANSIColors.RESET
+
+    def color_by_opname(self, opname: str) -> str:
+        if opname in (
+            "POP_TOP",
+            "POP_ITER",
+            "END_FOR",
+            "END_SEND",
+            "COPY",
+            "SWAP",
+            "PUSH_NULL",
+            "PUSH_EXC_INFO",
+            "NOP",
+            "CACHE",
+        ):
+            return self.op_stack
+
+        if opname.startswith(("LOAD_", "STORE_", "DELETE_", "IMPORT_")):
+            return self.op_load_store
+
+        if opname.startswith(("CALL", "RETURN")) or opname in (
+            "YIELD_VALUE",
+            "MAKE_FUNCTION",
+            "SET_FUNCTION_ATTRIBUTE",
+            "RESUME",
+        ):
+            return self.op_call_return
+
+        if opname.startswith(("BINARY_", "UNARY_")) or opname in (
+            "COMPARE_OP",
+            "IS_OP",
+            "CONTAINS_OP",
+            "GET_ITER",
+            "GET_YIELD_FROM_ITER",
+            "TO_BOOL",
+            "DELETE_SUBSCR",
+        ):
+            return self.op_binary_unary
+
+        if opname.startswith(("JUMP_", "POP_JUMP_", "FOR_ITER")) or opname in (
+            "SEND",
+            "GET_AWAITABLE",
+            "GET_AITER",
+            "GET_ANEXT",
+            "END_ASYNC_FOR",
+            "CLEANUP_THROW",
+        ):
+            return self.op_control_flow
+
+        if opname.startswith(
+            ("BUILD_", "LIST_", "DICT_", "UNPACK_")
+        ) or opname in ("SET_ADD", "MAP_ADD", "SET_UPDATE"):
+            return self.op_build
+
+        if opname.startswith(("SETUP_", "CHECK_")) or opname in (
+            "POP_EXCEPT",
+            "RERAISE",
+            "WITH_EXCEPT_START",
+            "RAISE_VARARGS",
+            "POP_BLOCK",
+        ):
+            return self.op_exceptions
+
+        return self.op_other
+
+
 @dataclass(frozen=True, kw_only=True)
 class LiveProfiler(ThemeSection):
     """Theme section for the live profiling TUI (Tachyon profiler).
@@ -343,7 +425,6 @@ class Unittest(ThemeSection):
     fail_info: str = ANSIColors.BOLD_RED
     reset: str = ANSIColors.RESET
 
-
 @dataclass(frozen=True, kw_only=True)
 class Theme:
     """A suite of themes for all sections of Python.
@@ -357,6 +438,7 @@ class Theme:
     syntax: Syntax = field(default_factory=Syntax)
     traceback: Traceback = field(default_factory=Traceback)
     unittest: Unittest = field(default_factory=Unittest)
+    dis: Dis = field(default_factory=Dis)
 
     def copy_with(
         self,
@@ -397,6 +479,7 @@ def no_colors(cls) -> Self:
             syntax=Syntax.no_colors(),
             traceback=Traceback.no_colors(),
             unittest=Unittest.no_colors(),
+            dis=Dis.no_colors(),
         )
 
 

Lib/dis.py

@@ -436,6 +436,9 @@ def __str__(self):
         formatter.print_instruction(self, False)
         return output.getvalue()
 
+def get_dis_theme():
+    from _colorize import get_theme
+    return get_theme().dis
 
 class Formatter:
 
@@ -478,8 +481,9 @@ def print_instruction(self, instr, mark_as_current=False):
                                     False, None, None, instr.positions),
                         False)
 
-    def print_instruction_line(self, instr, mark_as_current):
+    def print_instruction_line(self, instr, mark_as_current) -> None:
         """Format instruction details for inclusion in disassembly output."""
+        theme = get_dis_theme()
         lineno_width = self.lineno_width
         offset_width = self.offset_width
         label_width = self.label_width
@@ -527,7 +531,7 @@ def print_instruction_line(self, instr, mark_as_current):
         else:
             fields.append('   ')
         # Column: Opcode name
-        fields.append(instr.opname.ljust(_OPNAME_WIDTH))
+        fields.append(f"{theme.color_by_opname(instr.opname)}{instr.opname.ljust(_OPNAME_WIDTH)}{theme.reset}")
         # Column: Opcode argument
         if instr.arg is not None:
             # If opname is longer than _OPNAME_WIDTH, we allow it to overflow into
@@ -537,19 +541,20 @@ def print_instruction_line(self, instr, mark_as_current):
             fields.append(repr(instr.arg).rjust(_OPARG_WIDTH - opname_excess))
             # Column: Opcode argument details
             if instr.argrepr:
-                fields.append('(' + instr.argrepr + ')')
+                fields.append(f'{theme.argument_detail}(' + instr.argrepr + f'){theme.reset}')
         print(' '.join(fields).rstrip(), file=self.file)
 
     def print_exception_table(self, exception_entries):
         file = self.file
+        theme = get_dis_theme()
         if exception_entries:
             print("ExceptionTable:", file=file)
             for entry in exception_entries:
                 lasti = " lasti" if entry.lasti else ""
                 start = entry.start_label
                 end = entry.end_label
                 target = entry.target_label
-                print(f"  L{start} to L{end} -> L{target} [{entry.depth}]{lasti}", file=file)
+                print(f"  {theme.exception_label}L{start}{theme.reset} to {theme.exception_label}L{end}{theme.reset} -> {theme.exception_label}L{target}{theme.reset} [{entry.depth}]{lasti}", file=file)
 
 
 class ArgResolver:
@@ -833,13 +838,14 @@ def disassemble(co, lasti=-1, *, file=None, show_caches=False, adaptive=False,
 
 def _disassemble_recursive(co, *, file=None, depth=None, show_caches=False, adaptive=False, show_offsets=False, show_positions=False):
     disassemble(co, file=file, show_caches=show_caches, adaptive=adaptive, show_offsets=show_offsets, show_positions=show_positions)
+    theme = get_dis_theme()
     if depth is None or depth > 0:
         if depth is not None:
             depth = depth - 1
         for x in co.co_consts:
             if hasattr(x, 'co_code'):
                 print(file=file)
-                print("Disassembly of %r:" % (x,), file=file)
+                print(f"{theme.label_bg}{theme.label_fg}Disassembly of {x!r}:{theme.reset}", file=file)
                 _disassemble_recursive(
                     x, file=file, depth=depth, show_caches=show_caches,
                     adaptive=adaptive, show_offsets=show_offsets, show_positions=show_positions

Lib/test/test_compiler_assemble.py

@@ -4,7 +4,7 @@
 import types
 
 from test.support.bytecode_helper import AssemblerTestCase
-
+from test.support import force_not_colorized
 
 # Tests for the code-object creation stage of the compiler.
 
@@ -115,6 +115,7 @@ def inner():
         self.assemble_test(instructions, metadata, expected)
 
 
+    @force_not_colorized
     def test_exception_table(self):
         metadata = {
             'filename' : 'exc.py',

Lib/test/test_dis.py

@@ -13,9 +13,9 @@
 import textwrap
 import types
 import unittest
-from test.support import (captured_stdout, requires_debug_ranges,
-                          requires_specialization, cpython_only,
-                          os_helper, import_helper, reset_code)
+from test.support import (captured_stdout, force_not_colorized_test_class,
+                          requires_debug_ranges, requires_specialization,
+                          cpython_only, os_helper, import_helper, reset_code)
 from test.support.bytecode_helper import BytecodeTestCase
 
 
@@ -992,6 +992,7 @@ def do_disassembly_compare(self, got, expected):
         self.assertEqual(got, expected)
 
 
+@force_not_colorized_test_class
 class DisTests(DisTestBase):
 
     maxDiff = None
@@ -1468,6 +1469,7 @@ def f():
         self.assertEqual(assem_op, assem_cache)
 
 
+@force_not_colorized_test_class
 class DisWithFileTests(DisTests):
 
     # Run the tests again, using the file arg instead of print
@@ -1990,6 +1992,7 @@ def assertInstructionsEqual(self, instrs_1, instrs_2, /):
         instrs_2 = [instr_2._replace(positions=None, cache_info=None) for instr_2 in instrs_2]
         self.assertEqual(instrs_1, instrs_2)
 
+@force_not_colorized_test_class
 class InstructionTests(InstructionTestCase):
 
     def __init__(self, *args):
@@ -2311,6 +2314,7 @@ def test_cache_offset_and_end_offset(self):
 
 # get_instructions has its own tests above, so can rely on it to validate
 # the object oriented API
+@force_not_colorized_test_class
 class BytecodeTests(InstructionTestCase, DisTestBase):
 
     def test_instantiation(self):
@@ -2442,6 +2446,7 @@ def func():
         self.assertEqual(offsets, [0, 2])
 
 
+@force_not_colorized_test_class
 class TestDisTraceback(DisTestBase):
     def setUp(self) -> None:
         try:  # We need to clean up existing tracebacks
@@ -2479,6 +2484,7 @@ def test_distb_explicit_arg(self):
         self.do_disassembly_compare(self.get_disassembly(tb), dis_traceback)
 
 
+@force_not_colorized_test_class
 class TestDisTracebackWithFile(TestDisTraceback):
     # Run the `distb` tests again, using the file arg instead of print
     def get_disassembly(self, tb):
@@ -2513,6 +2519,7 @@ def _unroll_caches_as_Instructions(instrs, show_caches=False):
                                 False, None, None, instr.positions)
 
 
+@force_not_colorized_test_class
 class TestDisCLI(unittest.TestCase):
 
     def setUp(self):