diff --git a/Lib/pickletools.py b/Lib/pickletools.py index 53f25ea4e46b4d..5a50417c8217fe 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) annotate = 30 if args.annotate else 0 memo = {} if args.memo else None if args.output is None: @@ -2883,3 +2883,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..a26bb8a80d3bf3 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,174 @@ 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): + 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): + stderr = io.StringIO() + stdout = io.StringIO() + with ( + contextlib.redirect_stdout(stdout), + contextlib.redirect_stderr(stderr), + ): + pickletools._main(args=[*flags, self.filename]) + self.assertEqual(stderr.getvalue(), '') + return self.text_normalize(stdout.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() + self.addCleanup(os_helper.unlink, output_file) + 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) + + def test_unknown_flag(self): + with self.assertRaises(SystemExit): + 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',) + 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()) + + 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