gh-131178: Add tests for pickle command-line interface (#131275)

Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>
This commit is contained in:
Semyon Moroz 2025-04-06 19:50:32 +04:00 committed by GitHub
parent dbfc6a417a
commit f247e1d04c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 64 additions and 5 deletions

View file

@ -1907,7 +1907,7 @@ except ImportError:
dump, dumps, load, loads = _dump, _dumps, _load, _loads
if __name__ == "__main__":
def _main(args=None):
import argparse
import pprint
parser = argparse.ArgumentParser(
@ -1915,7 +1915,7 @@ if __name__ == "__main__":
parser.add_argument(
'pickle_file',
nargs='+', help='the pickle file')
args = parser.parse_args()
args = parser.parse_args(args)
for fn in args.pickle_file:
if fn == '-':
obj = load(sys.stdin.buffer)
@ -1923,3 +1923,7 @@ if __name__ == "__main__":
with open(fn, 'rb') as f:
obj = load(f)
pprint.pprint(obj)
if __name__ == "__main__":
_main()

View file

@ -1,18 +1,21 @@
from _compat_pickle import (IMPORT_MAPPING, REVERSE_IMPORT_MAPPING,
NAME_MAPPING, REVERSE_NAME_MAPPING)
import builtins
import pickle
import io
import collections
import contextlib
import io
import pickle
import struct
import sys
import tempfile
import warnings
import weakref
from textwrap import dedent
import doctest
import unittest
from test import support
from test.support import import_helper
from test.support import import_helper, os_helper
from test.pickletester import AbstractHookTests
from test.pickletester import AbstractUnpickleTests
@ -699,6 +702,58 @@ class CompatPickleTests(unittest.TestCase):
('multiprocessing.context', name))
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 dedent(string).strip()
def set_pickle_data(self, data):
with open(self.filename, 'wb') as f:
pickle.dump(data, f)
def invoke_pickle(self, *flags):
output = io.StringIO()
with contextlib.redirect_stdout(output):
pickle._main(args=[*flags, self.filename])
return self.text_normalize(output.getvalue())
def test_invocation(self):
# test 'python -m pickle pickle_file'
data = {
'a': [1, 2.0, 3+4j],
'b': ('character string', b'byte string'),
'c': 'string'
}
expect = '''
{'a': [1, 2.0, (3+4j)],
'b': ('character string', b'byte string'),
'c': 'string'}
'''
self.set_pickle_data(data)
with self.subTest(data=data):
res = self.invoke_pickle()
expect = self.text_normalize(expect)
self.assertListEqual(res.splitlines(), expect.splitlines())
def test_unknown_flag(self):
stderr = io.StringIO()
with self.assertRaises(SystemExit):
# check that the parser help is shown
with contextlib.redirect_stderr(stderr):
_ = self.invoke_pickle('--unknown')
self.assertStartsWith(stderr.getvalue(), 'usage: ')
def load_tests(loader, tests, pattern):
tests.addTest(doctest.DocTestSuite(pickle))
return tests