gh-149026: Add colour to `pickletools` CLI output (#149027)

Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>

diff --git a/Doc/library/pickletools.rst b/Doc/library/pickletools.rst
index 7a771ea3ab93d4146c216d2f1a6269311906ecad..e753ad3b08b81a1c2cbc8aa52aeb0c85e6a8f47f 100644 (file)
--- a/Doc/library/pickletools.rst
+++ b/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

diff --git a/Doc/whatsnew/3.15.rst b/Doc/whatsnew/3.15.rst
index eb08f8c4ed69e72b2d238f523267dc1a722a0df0..3c2c7a7e399d09c854ccf487872e951a2fe2ed02 100644 (file)
--- a/Doc/whatsnew/3.15.rst
+++ b/Doc/whatsnew/3.15.rst
@@ -1053,6 +1053,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
 ------
 

diff --git a/Lib/_colorize.py b/Lib/_colorize.py
index 379ca2529b6585d5376d1e9333b8d7c8ec88f478..62806b1d8d7bcf2dffce19dc349b6f03975ebe53 100644 (file)
--- a/Lib/_colorize.py
+++ b/Lib/_colorize.py
@@ -359,6 +359,23 @@ LiveProfilerLight = LiveProfiler(
 )
 
 
+@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 @@ class Theme:
         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 @@ class Theme:
             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 @@ class Theme:
             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(),

diff --git a/Lib/pickletools.py b/Lib/pickletools.py
index 29baf3be7ebb6e6c1ca46bfc0ce9a0aeefd0c3ef..976e218db192980bfa24e1fd49c7095a2dc9abb9 100644 (file)
--- a/Lib/pickletools.py
+++ b/Lib/pickletools.py
@@ -16,6 +16,8 @@ import pickle
 import re
 import sys
 
+lazy from _colorize import decolor, get_theme
+
 __all__ = ['dis', 'genops', 'optimize']
 
 bytes_types = pickle.bytes_types
@@ -2209,6 +2211,32 @@ for i, d in enumerate(opcodes):
     name2i[d.name] = i
     code2i[d.code] = i
 
+# 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
+})
+assert _opcode_color_attr.keys() <= name2i.keys(), (
+    f"unknown opcodes: {_opcode_color_attr.keys() - name2i.keys()}"
+)
 del name2i, code2i, i, d
 
 ##############################################################################
@@ -2443,13 +2471,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 +2544,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 +2583,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 +2887,7 @@ __test__ = {'disassembler_test': _dis_test,
 
 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')

diff --git a/Lib/test/test_pickletools.py b/Lib/test/test_pickletools.py
index 57285ddf6ebef5a661a08acfdaada86145a4d4ce..caf2d7ba6bfd8f55b3349b57a8864b1bef2f7add 100644 (file)
--- a/Lib/test/test_pickletools.py
+++ b/Lib/test/test_pickletools.py
@@ -160,6 +160,7 @@ class GenopsTests(unittest.TestCase):
             next(it)
 
 
+@support.force_not_colorized_test_class
 class DisTests(unittest.TestCase):
     maxDiff = None
 
@@ -518,6 +519,7 @@ class MiscTestCase(unittest.TestCase):
         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()

diff --git a/Misc/NEWS.d/next/Library/2026-04-26-23-01-50.gh-issue-149026.Akk4Bc.rst b/Misc/NEWS.d/next/Library/2026-04-26-23-01-50.gh-issue-149026.Akk4Bc.rst
new file mode 100644 (file)
index 0000000..d12a92e
--- /dev/null
+++ b/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.