https://github.com/python/cpython/commit/5c25bc5763ccaa0ee4bb6ebf7d6fa23e036197a1
commit: 5c25bc5763ccaa0ee4bb6ebf7d6fa23e036197a1
branch: main
author: Semyon Moroz <[email protected]>
committer: hugovk <[email protected]>
date: 2025-11-22T19:17:06+02:00
summary:
gh-131178: Add tests for `pickletools` command-line interface (#131287)
files:
M Lib/pickletools.py
M Lib/test/test_pickletools.py
diff --git a/Lib/pickletools.py b/Lib/pickletools.py
index 254b6c7fcc9dd2..29baf3be7ebb6e 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',
@@ -2864,7 +2864,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:
@@ -2885,3 +2885,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 cf990874621eae..57285ddf6ebef5 100644
--- a/Lib/test/test_pickletools.py
+++ b/Lib/test/test_pickletools.py
@@ -1,7 +1,11 @@
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 +518,170 @@ 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):
+ with (
+ support.captured_stdout() as stdout,
+ support.captured_stderr() as 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):
+ with support.captured_stderr() as stderr:
+ pickletools._main(args=['--unknown'])
+ 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 = 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
+ '''
+ 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 = 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
+ '''
+ 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 = 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
+ '''
+ 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 = 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
+ '''
+ 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 = 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)
+ 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)
+ with support.captured_stdout() as stdout:
+ pickletools._main(args=[flag, self.filename,
self.filename])
+ res = self.text_normalize(stdout.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
_______________________________________________
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]