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

Lib/_colorize.py

@@ -344,6 +344,88 @@ class Unittest(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 Theme:
     """A suite of themes for all sections of Python.
@@ -357,6 +439,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 +480,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: Instruction, mark_as_current: bool) -> 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

@@ -1,11 +1,14 @@
 import dis
 import io
+import os
 import textwrap
 import types
 
 from test.support.bytecode_helper import AssemblerTestCase
 
 
+os.environ.setdefault("NO_COLOR", "1")
+
 # Tests for the code-object creation stage of the compiler.
 
 class IsolatedAssembleTests(AssemblerTestCase):

Lib/test/test_dis.py

@@ -7,6 +7,7 @@
 import io
 import itertools
 import opcode
+import os
 import re
 import sys
 import tempfile
@@ -19,6 +20,8 @@
 from test.support.bytecode_helper import BytecodeTestCase
 
 
+os.environ.setdefault("NO_COLOR", "1")
+
 CACHE = dis.opmap["CACHE"]
 
 def get_tb():