From 83afce1f30b9539e05ca7ae557b126df880eed4f Mon Sep 17 00:00:00 2001 From: donBarbos Date: Sat, 15 Mar 2025 19:07:22 +0400 Subject: [PATCH 1/6] Add tests for `pickletools` command-line interface --- Lib/pickletools.py | 8 +- Lib/test/test_pickletools.py | 171 +++++++++++++++++++++++++++++++++++ 2 files changed, 177 insertions(+), 2 deletions(-) diff --git a/Lib/pickletools.py b/Lib/pickletools.py index 02aad12985dafe..fec7db3520a18d 100644 --- a/Lib/pickletools.py +++ b/Lib/pickletools.py @@ -2839,7 +2839,7 @@ def __init__(self, value): } -if __name__ == "__main__": +def _main(args=None): import argparse parser = argparse.ArgumentParser( description='disassemble one or more pickle files') @@ -2862,7 +2862,7 @@ def __init__(self, value): '-p', '--preamble', default="==> {name} <==", help='if more than one pickle file is specified, print this before' ' each disassembly') - args = parser.parse_args() + args = parser.parse_args(args) if not args.pickle_file: parser.print_help() else: @@ -2886,3 +2886,7 @@ def __init__(self, value): finally: if output is not sys.stdout: output.close() + + +if __name__ == "__main__": + _main() diff --git a/Lib/test/test_pickletools.py b/Lib/test/test_pickletools.py index a178d3353eecdf..be7913f90d9ed5 100644 --- a/Lib/test/test_pickletools.py +++ b/Lib/test/test_pickletools.py @@ -1,7 +1,12 @@ +import contextlib import io +import itertools import pickle import pickletools +import tempfile +import textwrap from test import support +from test.support import os_helper from test.pickletester import AbstractPickleTests import doctest import unittest @@ -514,6 +519,172 @@ def test__all__(self): support.check__all__(self, pickletools, not_exported=not_exported) +class CommandLineTest(unittest.TestCase): + def setUp(self): + self.filename = tempfile.mktemp() + self.addCleanup(os_helper.unlink, self.filename) + + @staticmethod + def text_normalize(string): + """Dedent *string* and strip it from its surrounding whitespaces. + + This method is used by the other utility functions so that any + string to write or to match against can be freely indented. + """ + return textwrap.dedent(string).strip() + + def set_pickle_data(self, data): + with open(self.filename, 'wb') as f: + pickle.dump(data, f) + + def invoke_pickletools(self, *flags): + output = io.StringIO() + with contextlib.redirect_stdout(output): + pickletools._main(args=[*flags, self.filename]) + return self.text_normalize(output.getvalue()) + + def check_output(self, data, expect, *flags): + with self.subTest(data=data, flags=flags): + self.set_pickle_data(data) + res = self.invoke_pickletools(*flags) + expect = self.text_normalize(expect) + self.assertListEqual(res.splitlines(), expect.splitlines()) + + def test_invocation(self): + # test various combinations of parameters + output_file = tempfile.mktemp() + base_flags = [ + (f'-o={output_file}', f'--output={output_file}'), + ('-m', '--memo'), + ('-l=2', '--indentlevel=2'), + ('-a', '--annotate'), + ('-p="Another:"', '--preamble="Another:"'), + ] + data = { "a", "b", "c" } + + self.set_pickle_data(data) + + for r in range(1, len(base_flags) + 1): + for choices in itertools.combinations(base_flags, r=r): + for args in itertools.product(*choices): + with self.subTest(args=args[1:]): + _ = self.invoke_pickletools(*args) + self.addCleanup(os_helper.unlink, output_file) + + with self.assertRaises(SystemExit): + # suppress argparse error message + with contextlib.redirect_stderr(io.StringIO()): + _ = self.invoke_pickletools('--unknown') + + def test_output_flag(self): + # test 'python -m pickletools -o/--output' + output_file = tempfile.mktemp() + data = ("fake_data",) + expect = ''' + 0: \\x80 PROTO 5 + 2: \\x95 FRAME 15 + 11: \\x8c SHORT_BINUNICODE 'fake_data' + 22: \\x94 MEMOIZE (as 0) + 23: \\x85 TUPLE1 + 24: \\x94 MEMOIZE (as 1) + 25: . STOP + highest protocol among opcodes = 4 + ''' + for flag in [f'-o={output_file}', f'--output={output_file}']: + with self.subTest(data=data, flags=flag): + self.set_pickle_data(data) + res = self.invoke_pickletools(flag) + with open(output_file, 'r') as f: + res_from_file = self.text_normalize(f.read()) + expect = self.text_normalize(expect) + + self.assertListEqual(res.splitlines(), []) + self.assertListEqual(res_from_file.splitlines(), + expect.splitlines()) + self.addCleanup(os_helper.unlink, output_file) + + def test_memo_flag(self): + # test 'python -m pickletools -m/--memo' + data = ("fake_data",) + expect = ''' + 0: \\x80 PROTO 5 + 2: \\x95 FRAME 15 + 11: \\x8c SHORT_BINUNICODE 'fake_data' + 22: \\x94 MEMOIZE (as 0) + 23: \\x85 TUPLE1 + 24: \\x94 MEMOIZE (as 1) + 25: . STOP + highest protocol among opcodes = 4 + ''' + for flag in ['-m', '--memo']: + self.check_output(data, expect, flag) + + def test_indentlevel_flag(self): + # test 'python -m pickletools -l/--indentlevel' + data = ("fake_data",) + expect = ''' + 0: \\x80 PROTO 5 + 2: \\x95 FRAME 15 + 11: \\x8c SHORT_BINUNICODE 'fake_data' + 22: \\x94 MEMOIZE (as 0) + 23: \\x85 TUPLE1 + 24: \\x94 MEMOIZE (as 1) + 25: . STOP + highest protocol among opcodes = 4 + ''' + for flag in ['-l=2', '--indentlevel=2']: + self.check_output(data, expect, flag) + + def test_annotate_flag(self): + # test 'python -m pickletools -a/--annotate' + data = ("fake_data",) + expect = ''' + 0: \\x80 PROTO 5 Protocol version indicator. + 2: \\x95 FRAME 15 Indicate the beginning of a new frame. + 11: \\x8c SHORT_BINUNICODE 'fake_data' Push a Python Unicode string object. + 22: \\x94 MEMOIZE (as 0) Store the stack top into the memo. The stack is not popped. + 23: \\x85 TUPLE1 Build a one-tuple out of the topmost item on the stack. + 24: \\x94 MEMOIZE (as 1) Store the stack top into the memo. The stack is not popped. + 25: . STOP Stop the unpickling machine. + highest protocol among opcodes = 4 + ''' + for flag in ['-a', '--annotate']: + self.check_output(data, expect, flag) + + def test_preamble_flag(self): + # test 'python -m pickletools -p/--preamble' + data = ("fake_data",) + expect = ''' + Another: + 0: \\x80 PROTO 5 + 2: \\x95 FRAME 15 + 11: \\x8c SHORT_BINUNICODE 'fake_data' + 22: \\x94 MEMOIZE (as 0) + 23: \\x85 TUPLE1 + 24: \\x94 MEMOIZE (as 1) + 25: . STOP + highest protocol among opcodes = 4 + Another: + 0: \\x80 PROTO 5 + 2: \\x95 FRAME 15 + 11: \\x8c SHORT_BINUNICODE 'fake_data' + 22: \\x94 MEMOIZE (as 0) + 23: \\x85 TUPLE1 + 24: \\x94 MEMOIZE (as 1) + 25: . STOP + highest protocol among opcodes = 4 + ''' + for flag in ['-p=Another:', '--preamble=Another:']: + with self.subTest(data=data, flags=flag): + self.set_pickle_data(data) + output = io.StringIO() + with contextlib.redirect_stdout(output): + pickletools._main(args=[flag, self.filename, self.filename]) + res = self.text_normalize(output.getvalue()) + expect = self.text_normalize(expect) + self.assertListEqual(res.splitlines(), expect.splitlines()) + + def load_tests(loader, tests, pattern): tests.addTest(doctest.DocTestSuite(pickletools)) return tests From 3292d1e954c0ea1bac511a542e40c1df838e9629 Mon Sep 17 00:00:00 2001 From: donBarbos Date: Sat, 15 Mar 2025 20:04:24 +0400 Subject: [PATCH 2/6] move addCleanup call to tempfile --- Lib/test/test_pickletools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Lib/test/test_pickletools.py b/Lib/test/test_pickletools.py index be7913f90d9ed5..38abdb57727f66 100644 --- a/Lib/test/test_pickletools.py +++ b/Lib/test/test_pickletools.py @@ -553,6 +553,7 @@ def check_output(self, data, expect, *flags): def test_invocation(self): # test various combinations of parameters output_file = tempfile.mktemp() + self.addCleanup(os_helper.unlink, output_file) base_flags = [ (f'-o={output_file}', f'--output={output_file}'), ('-m', '--memo'), @@ -569,7 +570,6 @@ def test_invocation(self): for args in itertools.product(*choices): with self.subTest(args=args[1:]): _ = self.invoke_pickletools(*args) - self.addCleanup(os_helper.unlink, output_file) with self.assertRaises(SystemExit): # suppress argparse error message @@ -579,6 +579,7 @@ def test_invocation(self): def test_output_flag(self): # test 'python -m pickletools -o/--output' output_file = tempfile.mktemp() + self.addCleanup(os_helper.unlink, output_file) data = ("fake_data",) expect = ''' 0: \\x80 PROTO 5 @@ -601,7 +602,6 @@ def test_output_flag(self): self.assertListEqual(res.splitlines(), []) self.assertListEqual(res_from_file.splitlines(), expect.splitlines()) - self.addCleanup(os_helper.unlink, output_file) def test_memo_flag(self): # test 'python -m pickletools -m/--memo' From 2d9d6447d2e401af908bde84ed862dd35d6b088d Mon Sep 17 00:00:00 2001 From: donBarbos Date: Thu, 27 Mar 2025 18:41:50 +0400 Subject: [PATCH 3/6] Resolve conflicts --- Lib/pickletools.py | 41 +++++++++++++++++++---------------------- 1 file changed, 19 insertions(+), 22 deletions(-) diff --git a/Lib/pickletools.py b/Lib/pickletools.py index fec7db3520a18d..121f00d65d809e 100644 --- a/Lib/pickletools.py +++ b/Lib/pickletools.py @@ -2863,29 +2863,26 @@ def _main(args=None): help='if more than one pickle file is specified, print this before' ' each disassembly') args = parser.parse_args(args) - if not args.pickle_file: - parser.print_help() + annotate = 30 if args.annotate else 0 + memo = {} if args.memo else None + if args.output is None: + output = sys.stdout else: - annotate = 30 if args.annotate else 0 - memo = {} if args.memo else None - if args.output is None: - output = sys.stdout - else: - output = open(args.output, 'w') - try: - for arg in args.pickle_file: - if len(args.pickle_file) > 1: - name = '' if arg == '-' else arg - preamble = args.preamble.format(name=name) - output.write(preamble + '\n') - if arg == '-': - dis(sys.stdin.buffer, output, memo, args.indentlevel, annotate) - else: - with open(arg, 'rb') as f: - dis(f, output, memo, args.indentlevel, annotate) - finally: - if output is not sys.stdout: - output.close() + output = open(args.output, 'w') + try: + for arg in args.pickle_file: + if len(args.pickle_file) > 1: + name = '' if arg == '-' else arg + preamble = args.preamble.format(name=name) + output.write(preamble + '\n') + if arg == '-': + dis(sys.stdin.buffer, output, memo, args.indentlevel, annotate) + else: + with open(arg, 'rb') as f: + dis(f, output, memo, args.indentlevel, annotate) + finally: + if output is not sys.stdout: + output.close() if __name__ == "__main__": From 393e431564049f1fa0e34e3030f1445bc42020c0 Mon Sep 17 00:00:00 2001 From: donBarbos Date: Thu, 27 Mar 2025 18:47:57 +0400 Subject: [PATCH 4/6] Back end of file --- Lib/pickletools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Lib/pickletools.py b/Lib/pickletools.py index f0ef4330622d21..5a50417c8217fe 100644 --- a/Lib/pickletools.py +++ b/Lib/pickletools.py @@ -2886,4 +2886,4 @@ def _main(args=None): if __name__ == "__main__": - _main() \ No newline at end of file + _main() From ce3683917cc3224f7fa065e77093a48c538aeebb Mon Sep 17 00:00:00 2001 From: donBarbos Date: Tue, 20 May 2025 09:26:28 +0400 Subject: [PATCH 5/6] Refactor --- Lib/test/test_pickletools.py | 38 +++++++++++++++++++----------------- 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/Lib/test/test_pickletools.py b/Lib/test/test_pickletools.py index 38abdb57727f66..a26bb8a80d3bf3 100644 --- a/Lib/test/test_pickletools.py +++ b/Lib/test/test_pickletools.py @@ -526,11 +526,6 @@ def setUp(self): @staticmethod def text_normalize(string): - """Dedent *string* and strip it from its surrounding whitespaces. - - This method is used by the other utility functions so that any - string to write or to match against can be freely indented. - """ return textwrap.dedent(string).strip() def set_pickle_data(self, data): @@ -538,10 +533,15 @@ def set_pickle_data(self, data): pickle.dump(data, f) def invoke_pickletools(self, *flags): - output = io.StringIO() - with contextlib.redirect_stdout(output): + stderr = io.StringIO() + stdout = io.StringIO() + with ( + contextlib.redirect_stdout(stdout), + contextlib.redirect_stderr(stderr), + ): pickletools._main(args=[*flags, self.filename]) - return self.text_normalize(output.getvalue()) + self.assertEqual(stderr.getvalue(), '') + return self.text_normalize(stdout.getvalue()) def check_output(self, data, expect, *flags): with self.subTest(data=data, flags=flags): @@ -561,7 +561,7 @@ def test_invocation(self): ('-a', '--annotate'), ('-p="Another:"', '--preamble="Another:"'), ] - data = { "a", "b", "c" } + data = { 'a', 'b', 'c' } self.set_pickle_data(data) @@ -569,18 +569,20 @@ def test_invocation(self): for choices in itertools.combinations(base_flags, r=r): for args in itertools.product(*choices): with self.subTest(args=args[1:]): - _ = self.invoke_pickletools(*args) + self.invoke_pickletools(*args) + def test_unknown_flag(self): with self.assertRaises(SystemExit): - # suppress argparse error message - with contextlib.redirect_stderr(io.StringIO()): - _ = self.invoke_pickletools('--unknown') + output = io.StringIO() + with contextlib.redirect_stderr(output): + pickletools._main(args=['--unknown']) + self.assertStartsWith(output.getvalue(), 'usage: ') def test_output_flag(self): # test 'python -m pickletools -o/--output' output_file = tempfile.mktemp() self.addCleanup(os_helper.unlink, output_file) - data = ("fake_data",) + data = ('fake_data',) expect = ''' 0: \\x80 PROTO 5 2: \\x95 FRAME 15 @@ -605,7 +607,7 @@ def test_output_flag(self): def test_memo_flag(self): # test 'python -m pickletools -m/--memo' - data = ("fake_data",) + data = ('fake_data',) expect = ''' 0: \\x80 PROTO 5 2: \\x95 FRAME 15 @@ -621,7 +623,7 @@ def test_memo_flag(self): def test_indentlevel_flag(self): # test 'python -m pickletools -l/--indentlevel' - data = ("fake_data",) + data = ('fake_data',) expect = ''' 0: \\x80 PROTO 5 2: \\x95 FRAME 15 @@ -637,7 +639,7 @@ def test_indentlevel_flag(self): def test_annotate_flag(self): # test 'python -m pickletools -a/--annotate' - data = ("fake_data",) + data = ('fake_data',) expect = ''' 0: \\x80 PROTO 5 Protocol version indicator. 2: \\x95 FRAME 15 Indicate the beginning of a new frame. @@ -653,7 +655,7 @@ def test_annotate_flag(self): def test_preamble_flag(self): # test 'python -m pickletools -p/--preamble' - data = ("fake_data",) + data = ('fake_data',) expect = ''' Another: 0: \\x80 PROTO 5 From 1546d0e47f17fe4c96bd8c5766819535ab9324a4 Mon Sep 17 00:00:00 2001 From: donBarbos Date: Sat, 2 Aug 2025 15:08:33 +0400 Subject: [PATCH 6/6] Accept suggestions --- Lib/test/test_pickletools.py | 99 +++++++++++++++++------------------- 1 file changed, 47 insertions(+), 52 deletions(-) diff --git a/Lib/test/test_pickletools.py b/Lib/test/test_pickletools.py index a26bb8a80d3bf3..caf6d084603151 100644 --- a/Lib/test/test_pickletools.py +++ b/Lib/test/test_pickletools.py @@ -1,4 +1,3 @@ -import contextlib import io import itertools import pickle @@ -533,11 +532,9 @@ def set_pickle_data(self, data): pickle.dump(data, f) def invoke_pickletools(self, *flags): - stderr = io.StringIO() - stdout = io.StringIO() with ( - contextlib.redirect_stdout(stdout), - contextlib.redirect_stderr(stderr), + support.captured_stdout() as stdout, + support.captured_stderr() as stderr, ): pickletools._main(args=[*flags, self.filename]) self.assertEqual(stderr.getvalue(), '') @@ -573,23 +570,22 @@ def test_invocation(self): def test_unknown_flag(self): with self.assertRaises(SystemExit): - output = io.StringIO() - with contextlib.redirect_stderr(output): + with support.captured_stderr() as stderr: pickletools._main(args=['--unknown']) - self.assertStartsWith(output.getvalue(), 'usage: ') + self.assertStartsWith(stderr.getvalue(), 'usage: ') def test_output_flag(self): # test 'python -m pickletools -o/--output' output_file = tempfile.mktemp() self.addCleanup(os_helper.unlink, output_file) data = ('fake_data',) - expect = ''' - 0: \\x80 PROTO 5 - 2: \\x95 FRAME 15 - 11: \\x8c SHORT_BINUNICODE 'fake_data' - 22: \\x94 MEMOIZE (as 0) - 23: \\x85 TUPLE1 - 24: \\x94 MEMOIZE (as 1) + expect = r''' + 0: \x80 PROTO 5 + 2: \x95 FRAME 15 + 11: \x8c SHORT_BINUNICODE 'fake_data' + 22: \x94 MEMOIZE (as 0) + 23: \x85 TUPLE1 + 24: \x94 MEMOIZE (as 1) 25: . STOP highest protocol among opcodes = 4 ''' @@ -608,13 +604,13 @@ def test_output_flag(self): def test_memo_flag(self): # test 'python -m pickletools -m/--memo' data = ('fake_data',) - expect = ''' - 0: \\x80 PROTO 5 - 2: \\x95 FRAME 15 - 11: \\x8c SHORT_BINUNICODE 'fake_data' - 22: \\x94 MEMOIZE (as 0) - 23: \\x85 TUPLE1 - 24: \\x94 MEMOIZE (as 1) + expect = r''' + 0: \x80 PROTO 5 + 2: \x95 FRAME 15 + 11: \x8c SHORT_BINUNICODE 'fake_data' + 22: \x94 MEMOIZE (as 0) + 23: \x85 TUPLE1 + 24: \x94 MEMOIZE (as 1) 25: . STOP highest protocol among opcodes = 4 ''' @@ -624,13 +620,13 @@ def test_memo_flag(self): def test_indentlevel_flag(self): # test 'python -m pickletools -l/--indentlevel' data = ('fake_data',) - expect = ''' - 0: \\x80 PROTO 5 - 2: \\x95 FRAME 15 - 11: \\x8c SHORT_BINUNICODE 'fake_data' - 22: \\x94 MEMOIZE (as 0) - 23: \\x85 TUPLE1 - 24: \\x94 MEMOIZE (as 1) + expect = r''' + 0: \x80 PROTO 5 + 2: \x95 FRAME 15 + 11: \x8c SHORT_BINUNICODE 'fake_data' + 22: \x94 MEMOIZE (as 0) + 23: \x85 TUPLE1 + 24: \x94 MEMOIZE (as 1) 25: . STOP highest protocol among opcodes = 4 ''' @@ -640,13 +636,13 @@ def test_indentlevel_flag(self): def test_annotate_flag(self): # test 'python -m pickletools -a/--annotate' data = ('fake_data',) - expect = ''' - 0: \\x80 PROTO 5 Protocol version indicator. - 2: \\x95 FRAME 15 Indicate the beginning of a new frame. - 11: \\x8c SHORT_BINUNICODE 'fake_data' Push a Python Unicode string object. - 22: \\x94 MEMOIZE (as 0) Store the stack top into the memo. The stack is not popped. - 23: \\x85 TUPLE1 Build a one-tuple out of the topmost item on the stack. - 24: \\x94 MEMOIZE (as 1) Store the stack top into the memo. The stack is not popped. + expect = r''' + 0: \x80 PROTO 5 Protocol version indicator. + 2: \x95 FRAME 15 Indicate the beginning of a new frame. + 11: \x8c SHORT_BINUNICODE 'fake_data' Push a Python Unicode string object. + 22: \x94 MEMOIZE (as 0) Store the stack top into the memo. The stack is not popped. + 23: \x85 TUPLE1 Build a one-tuple out of the topmost item on the stack. + 24: \x94 MEMOIZE (as 1) Store the stack top into the memo. The stack is not popped. 25: . STOP Stop the unpickling machine. highest protocol among opcodes = 4 ''' @@ -656,33 +652,32 @@ def test_annotate_flag(self): def test_preamble_flag(self): # test 'python -m pickletools -p/--preamble' data = ('fake_data',) - expect = ''' + expect = r''' Another: - 0: \\x80 PROTO 5 - 2: \\x95 FRAME 15 - 11: \\x8c SHORT_BINUNICODE 'fake_data' - 22: \\x94 MEMOIZE (as 0) - 23: \\x85 TUPLE1 - 24: \\x94 MEMOIZE (as 1) + 0: \x80 PROTO 5 + 2: \x95 FRAME 15 + 11: \x8c SHORT_BINUNICODE 'fake_data' + 22: \x94 MEMOIZE (as 0) + 23: \x85 TUPLE1 + 24: \x94 MEMOIZE (as 1) 25: . STOP highest protocol among opcodes = 4 Another: - 0: \\x80 PROTO 5 - 2: \\x95 FRAME 15 - 11: \\x8c SHORT_BINUNICODE 'fake_data' - 22: \\x94 MEMOIZE (as 0) - 23: \\x85 TUPLE1 - 24: \\x94 MEMOIZE (as 1) + 0: \x80 PROTO 5 + 2: \x95 FRAME 15 + 11: \x8c SHORT_BINUNICODE 'fake_data' + 22: \x94 MEMOIZE (as 0) + 23: \x85 TUPLE1 + 24: \x94 MEMOIZE (as 1) 25: . STOP highest protocol among opcodes = 4 ''' for flag in ['-p=Another:', '--preamble=Another:']: with self.subTest(data=data, flags=flag): self.set_pickle_data(data) - output = io.StringIO() - with contextlib.redirect_stdout(output): + with support.captured_stdout() as stdout: pickletools._main(args=[flag, self.filename, self.filename]) - res = self.text_normalize(output.getvalue()) + res = self.text_normalize(stdout.getvalue()) expect = self.text_normalize(expect) self.assertListEqual(res.splitlines(), expect.splitlines())