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

Doc/whatsnew/3.15.rst

@@ -737,6 +737,15 @@ difflib
   (Contributed by Jiahao Li in :gh:`134580`.)
 
 
+dis
+---
+
+* :func:`dis.dis` supports colored output by default, which can also be
+  :ref:`controlled <using-on-controlling-color>` through ``NO_COLOR=1``
+  environment variable.
+  (Contributed by Abduaziz Ziyodov in :gh:`144207`.)
+
+
 functools
 ---------
 

Lib/_colorize.py

@@ -200,6 +200,50 @@ class Difflib(ThemeSection):
     reset: str = ANSIColors.RESET
 
 
+@dataclass(frozen=True, kw_only=True)
+class Dis(ThemeSection):
+    label_bg: str = ANSIColors.BACKGROUND_CYAN
+    label_fg: str = ANSIColors.BLACK
+
+    exception_label: str = ANSIColors.CYAN
+    argument_detail: str = ANSIColors.CYAN
+
+    op_load: str = ANSIColors.BOLD_BLUE
+    op_pop: str = ANSIColors.BOLD_MAGENTA
+    op_call_return: str = ANSIColors.BOLD_YELLOW
+    op_control_flow: str = ANSIColors.BOLD_GREEN
+
+    reset: str = ANSIColors.RESET
+
+    def color_by_opname(self, opname: str) -> str:
+        if opname.startswith("LOAD_"):
+            return self.op_load
+
+        if opname.startswith("POP_"):
+            return self.op_pop
+
+        if opname.startswith(("CALL", "RETURN")) or opname in (
+            "YIELD_VALUE",
+            "MAKE_FUNCTION",
+            "SET_FUNCTION_ATTRIBUTE",
+            "RESUME",
+        ):
+            return self.op_call_return
+
+        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
+
+
+        return self.reset
+
+
 @dataclass(frozen=True, kw_only=True)
 class LiveProfiler(ThemeSection):
     """Theme section for the live profiling TUI (Tachyon profiler).
@@ -357,6 +401,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,
@@ -367,6 +412,7 @@ def copy_with(
         syntax: Syntax | None = None,
         traceback: Traceback | None = None,
         unittest: Unittest | None = None,
+        dis: Dis | None = None
     ) -> Self:
         """Return a new Theme based on this instance with some sections replaced.
 
@@ -380,6 +426,7 @@ def copy_with(
             syntax=syntax or self.syntax,
             traceback=traceback or self.traceback,
             unittest=unittest or self.unittest,
+            dis=dis or self.dis
         )
 
     @classmethod
@@ -397,6 +444,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

@@ -437,6 +437,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:
 
@@ -481,6 +484,7 @@ def print_instruction(self, instr, mark_as_current=False):
 
     def print_instruction_line(self, instr, mark_as_current):
         """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
@@ -528,7 +532,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
@@ -538,19 +542,25 @@ 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 "
+                    f"{theme.exception_label}L{end}{theme.reset} "
+                    f"-> {theme.exception_label}L{target}{theme.reset} [{entry.depth}]{lasti}",
+                    file=file,
+                )
 
 
 class ArgResolver:
@@ -840,13 +850,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,10 @@
 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,
+                          force_colorized_test_class, requires_debug_ranges,
+                          requires_specialization, cpython_only, os_helper,
+                          import_helper, reset_code)
 from test.support.bytecode_helper import BytecodeTestCase
 
 
@@ -36,6 +37,8 @@ def _error():
 
 TRACEBACK_CODE = get_tb().tb_frame.f_code
 
+theme = dis._get_dis_theme()
+
 class _C:
     def __init__(self, x):
         self.x = x == 1
@@ -992,6 +995,7 @@ def do_disassembly_compare(self, got, expected):
         self.assertEqual(got, expected)
 
 
+@force_not_colorized_test_class
 class DisTests(DisTestBase):
 
     maxDiff = None
@@ -1468,6 +1472,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 +1995,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 +2317,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 +2449,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 +2487,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 +2522,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):
@@ -2626,6 +2636,65 @@ def test_specialized_code(self):
         for flag in ['-S', '--specialized']:
             self.check_output(source, expect, flag)
 
+@force_colorized_test_class
+class DisColoredTests(unittest.TestCase):
+    def get_colored_output(self, func):
+        output = io.StringIO()
+
+        with contextlib.redirect_stdout(output):
+            dis.dis(func)
+
+        return output.getvalue()
+
+    def assertOpColored(self, output, opname, color):
+        self.assertIn(
+            f"{color}{opname}", output,
+            f"{opname} should be colored with {color!r}"
+        )
+
+    def test_load_ops_colored(self):
+        def f(a):
+            return a
+        out = self.get_colored_output(f)
+        self.assertOpColored(out, "LOAD_FAST", theme.op_load)
+
+    def test_call_return_ops_colored(self):
+        def f():
+            return 1
+        out = self.get_colored_output(f)
+        self.assertOpColored(out, "RETURN_VALUE", theme.op_call_return)
+        self.assertOpColored(out, "RESUME", theme.op_call_return)
+
+    def test_pop_ops_colored(self):
+        def f(a):
+            print(a)
+        out = self.get_colored_output(f)
+        self.assertOpColored(out, "POP_TOP", theme.op_pop)
+
+    def test_control_flow_ops_colored(self):
+        def f(a):
+            for _ in a:
+                pass
+        out = self.get_colored_output(f)
+        self.assertOpColored(out, "FOR_ITER", theme.op_control_flow)
+        self.assertOpColored(out, "JUMP_BACKWARD", theme.op_control_flow)
+
+    def test_argrepr_colored(self):
+        def f(a):
+            print(a)
+        out = self.get_colored_output(f)
+        self.assertIn(f"{theme.argument_detail}(", out)
+
+    def test_color_by_opname_coverage(self):
+        self.assertEqual(theme.color_by_opname("LOAD_FAST"), theme.op_load)
+        self.assertEqual(theme.color_by_opname("LOAD_GLOBAL"), theme.op_load)
+        self.assertEqual(theme.color_by_opname("POP_TOP"), theme.op_pop)
+        self.assertEqual(theme.color_by_opname("CALL"), theme.op_call_return)
+        self.assertEqual(theme.color_by_opname("RETURN_VALUE"), theme.op_call_return)
+        self.assertEqual(theme.color_by_opname("RESUME"), theme.op_call_return)
+        self.assertEqual(theme.color_by_opname("FOR_ITER"), theme.op_control_flow)
+        self.assertEqual(theme.color_by_opname("JUMP_BACKWARD"), theme.op_control_flow)
+        self.assertEqual(theme.color_by_opname("BINARY_OP"), theme.reset)  # uncolored
 
 if __name__ == "__main__":
     unittest.main()

Misc/NEWS.d/next/Library/2026-01-25-15-26-23.gh-issue-144207.G2c_qd.rst

@@ -0,0 +1,3 @@
+:func:`dis.dis` supports colored output by default which can also be
+:ref:`controlled <using-on-controlling-color>` through ``NO_COLOR=1``
+environment variable. Contributed by Abduaziz Ziyodov.