https://github.com/python/cpython/commit/8851a06e6e7ff23d91e5568df03a5ade570d4ab5 commit: 8851a06e6e7ff23d91e5568df03a5ade570d4ab5 branch: main author: Hugo van Kemenade <[email protected]> committer: hugovk <[email protected]> date: 2026-04-29T18:33:05+03:00 summary:
gh-149026: Add colour to `pickletools` CLI output (#149027) Co-authored-by: Bénédikt Tran <[email protected]> files: A Misc/NEWS.d/next/Library/2026-04-26-23-01-50.gh-issue-149026.Akk4Bc.rst M Doc/library/pickletools.rst M Doc/whatsnew/3.15.rst M Lib/_colorize.py M Lib/pickletools.py M Lib/test/test_pickletools.py diff --git a/Doc/library/pickletools.rst b/Doc/library/pickletools.rst index 7a771ea3ab93d4..e753ad3b08b81a 100644 --- 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 eb08f8c4ed69e7..3c2c7a7e399d09 100644 --- 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 379ca2529b6585..62806b1d8d7bcf 100644 --- a/Lib/_colorize.py +++ b/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(), diff --git a/Lib/pickletools.py b/Lib/pickletools.py index 29baf3be7ebb6e..976e218db19298 100644 --- a/Lib/pickletools.py +++ b/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 @@ -2209,6 +2211,32 @@ def __init__(self, name, code, arg, 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 @@ 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') diff --git a/Lib/test/test_pickletools.py b/Lib/test/test_pickletools.py index 57285ddf6ebef5..caf2d7ba6bfd8f 100644 --- a/Lib/test/test_pickletools.py +++ b/Lib/test/test_pickletools.py @@ -160,6 +160,7 @@ def test_unknown_opcode_without_pos(self): next(it) [email protected]_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) [email protected]_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 index 00000000000000..d12a92e9f530da --- /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. _______________________________________________ Python-checkins mailing list -- [email protected] To unsubscribe send an email to [email protected] https://mail.python.org/mailman3//lists/python-checkins.python.org Member address: [email protected]
