gh-149026: Add colour to pickletools CLI output

Doc/library/pickletools.rst

@@ -79,6 +79,9 @@ Command-line options
 
    A pickle file to read, or ``-`` to indicate reading from standard input.
 
+.. versionadded:: next
+   Output is in color by default and can be
+   :ref:`controlled using environment variables <using-on-controlling-color>`.
 
 
 Programmatic interface

Doc/whatsnew/3.15.rst

@@ -1010,6 +1010,15 @@ pickle
   (Contributed by Zackery Spytz and Serhiy Storchaka in :gh:`77188`.)
 
 
+pickletools
+-----------
+
+* The output of the :mod:`pickletools` command-line interface is colored by
+  default. This can be controlled with
+  :ref:`environment variables <using-on-controlling-color>`.
+  (Contributed by Hugo van Kemenade in :gh:`149026`.)
+
+
 pprint
 ------
 

Lib/_colorize.py

@@ -359,6 +359,23 @@ class LiveProfiler(ThemeSection):
 )
 
 
+@dataclass(frozen=True, kw_only=True)
+class Pickletools(ThemeSection):
+    annotation: str = ANSIColors.GREY
+    arg_number: str = ANSIColors.YELLOW
+    arg_string: str = ANSIColors.GREEN
+    mark: str = ANSIColors.GREY
+    op_call: str = ANSIColors.GREEN
+    op_container: str = ANSIColors.INTENSE_BLUE
+    op_memo: str = ANSIColors.MAGENTA
+    op_meta: str = ANSIColors.GREY
+    op_stack: str = ANSIColors.BOLD_RED
+    opcode_code: str = ANSIColors.CYAN
+    position: str = ANSIColors.GREY
+    proto: str = ANSIColors.YELLOW
+    reset: str = ANSIColors.RESET
+
+
 @dataclass(frozen=True, kw_only=True)
 class Syntax(ThemeSection):
     prompt: str = ANSIColors.BOLD_MAGENTA
@@ -429,6 +446,7 @@ class Theme:
     fancycompleter: FancyCompleter = field(default_factory=FancyCompleter)
     http_server: HttpServer = field(default_factory=HttpServer)
     live_profiler: LiveProfiler = field(default_factory=LiveProfiler)
+    pickletools: Pickletools = field(default_factory=Pickletools)
     syntax: Syntax = field(default_factory=Syntax)
     timeit: Timeit = field(default_factory=Timeit)
     tokenize: Tokenize = field(default_factory=Tokenize)
@@ -444,6 +462,7 @@ def copy_with(
         fancycompleter: FancyCompleter | None = None,
         http_server: HttpServer | None = None,
         live_profiler: LiveProfiler | None = None,
+        pickletools: Pickletools | None = None,
         syntax: Syntax | None = None,
         timeit: Timeit | None = None,
         tokenize: Tokenize | None = None,
@@ -462,6 +481,7 @@ def copy_with(
             fancycompleter=fancycompleter or self.fancycompleter,
             http_server=http_server or self.http_server,
             live_profiler=live_profiler or self.live_profiler,
+            pickletools=pickletools or self.pickletools,
             syntax=syntax or self.syntax,
             timeit=timeit or self.timeit,
             tokenize=tokenize or self.tokenize,
@@ -484,6 +504,7 @@ def no_colors(cls) -> Self:
             fancycompleter=FancyCompleter.no_colors(),
             http_server=HttpServer.no_colors(),
             live_profiler=LiveProfiler.no_colors(),
+            pickletools=Pickletools.no_colors(),
             syntax=Syntax.no_colors(),
             timeit=Timeit.no_colors(),
             tokenize=Tokenize.no_colors(),

Lib/pickletools.py

@@ -16,6 +16,8 @@
 import re
 import sys
 
+lazy from _colorize import decolor, get_theme
+
 __all__ = ['dis', 'genops', 'optimize']
 
 bytes_types = pickle.bytes_types
@@ -2392,6 +2394,31 @@ def optimize(p):
 ##############################################################################
 # A symbolic pickle disassembler.
 
+# Group opcode names into categories for colourised CLI output.
+_opcode_categories = frozendict(
+    op_call=frozenset({
+        "BUILD", "EXT1", "EXT2", "EXT4", "GLOBAL", "INST", "NEWOBJ",
+        "NEWOBJ_EX", "OBJ", "REDUCE", "STACK_GLOBAL",
+    }),
+    op_container=frozenset({
+        "ADDITEMS", "APPEND", "APPENDS", "DICT", "EMPTY_DICT", "EMPTY_LIST",
+        "EMPTY_SET", "EMPTY_TUPLE", "FROZENSET", "LIST", "SETITEM",
+        "SETITEMS", "TUPLE", "TUPLE1", "TUPLE2", "TUPLE3",
+    }),
+    op_memo=frozenset({
+        "BINGET", "BINPUT", "GET", "LONG_BINGET", "LONG_BINPUT", "MEMOIZE",
+        "PUT",
+    }),
+    op_meta=frozenset({"BINPERSID", "FRAME", "MARK", "PERSID", "PROTO"}),
+    op_stack=frozenset({"DUP", "POP", "POP_MARK", "STOP"}),
+)
+_opcode_color_attr = frozendict({
+    name: attr
+    for attr, names in _opcode_categories.items()
+    for name in names
+})
+
+
 def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
     """Produce a symbolic disassembly of a pickle.
 
@@ -2443,13 +2470,19 @@ def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
     indentchunk = ' ' * indentlevel
     errormsg = None
     annocol = annotate  # column hint for annotations
+    t = get_theme(tty_file=out).pickletools
     for opcode, arg, pos in genops(pickle):
         if pos is not None:
-            print("%5d:" % pos, end=' ', file=out)
+            print(f"{t.position}{pos:5d}:{t.reset}", end=' ', file=out)
 
-        line = "%-4s %s%s" % (repr(opcode.code)[1:-1],
-                              indentchunk * len(markstack),
-                              opcode.name)
+        attr = _opcode_color_attr.get(opcode.name)
+        opcode_color = getattr(t, attr) if attr else ""
+        opcode_reset = t.reset if attr else ""
+        line = (
+            f"{t.opcode_code}{repr(opcode.code)[1:-1]:<4}{t.reset} "
+            f"{indentchunk * len(markstack)}"
+            f"{opcode_color}{opcode.name}{opcode_reset}"
+        )
 
         maxproto = max(maxproto, opcode.proto)
         before = opcode.stack_before    # don't mutate
@@ -2510,18 +2543,26 @@ def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
             line += ' ' * (10 - len(opcode.name))
             if arg is not None:
                 if opcode.name in ("STRING", "BINSTRING", "SHORT_BINSTRING"):
-                    line += ' ' + ascii(arg)
+                    arg_text = ascii(arg)
                 else:
-                    line += ' ' + repr(arg)
+                    arg_text = repr(arg)
+                arg_color = (
+                    t.arg_number
+                    if isinstance(arg, (int, float))
+                    else t.arg_string
+                )
+                line += f" {arg_color}{arg_text}{t.reset}"
             if markmsg:
-                line += ' ' + markmsg
+                line += f" {t.mark}{markmsg}{t.reset}"
         if annotate:
-            line += ' ' * (annocol - len(line))
+            visible_len = len(decolor(line))
+            line += ' ' * (annocol - visible_len)
             # make a mild effort to align annotations
-            annocol = len(line)
+            annocol = max(visible_len, annocol)
             if annocol > 50:
                 annocol = annotate
-            line += ' ' + opcode.doc.split('\n', 1)[0]
+            doc = opcode.doc.split('\n', 1)[0]
+            line += f" {t.annotation}{doc}{t.reset}"
         print(line, file=out)
 
         if errormsg:
@@ -2541,7 +2582,11 @@ def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
 
         stack.extend(after)
 
-    print("highest protocol among opcodes =", maxproto, file=out)
+    print(
+        "highest protocol among opcodes =",
+        f"{t.proto}{maxproto}{t.reset}",
+        file=out,
+    )
     if stack:
         raise ValueError("stack not empty after STOP: %r" % stack)
 
@@ -2841,10 +2886,7 @@ def __init__(self, value):
 
 def _main(args=None):
     import argparse
-    parser = argparse.ArgumentParser(
-        description='disassemble one or more pickle files',
-        color=True,
-    )
+    parser = argparse.ArgumentParser(description='disassemble one or more pickle files')
     parser.add_argument(
         'pickle_file',
         nargs='+', help='the pickle file')

Lib/test/test_pickletools.py

@@ -160,6 +160,7 @@ def test_unknown_opcode_without_pos(self):
             next(it)
 
 
+@support.force_not_colorized_test_class
 class DisTests(unittest.TestCase):
     maxDiff = None
 
@@ -518,6 +519,7 @@ def test__all__(self):
         support.check__all__(self, pickletools, not_exported=not_exported)
 
 
+@support.force_not_colorized_test_class
 class CommandLineTest(unittest.TestCase):
     def setUp(self):
         self.filename = tempfile.mktemp()

Misc/NEWS.d/next/Library/2026-04-26-23-01-50.gh-issue-149026.Akk4Bc.rst

@@ -0,0 +1 @@
+Add colour to :mod:`pickletools` CLI output. Patch by Hugo van Kemenade.