Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Apr 6, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions Lib/pickle.py
Original file line number Diff line number Diff line change
Expand Up @@ -1907,19 +1907,23 @@ def _loads(s, /, *, fix_imports=True, encoding="ASCII", errors="strict",
dump, dumps, load, loads = _dump, _dumps, _load, _loads


if __name__ == "__main__":
def _main(args=None):
import argparse
import pprint
parser = argparse.ArgumentParser(
description='display contents of the pickle files')
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)
else:
with open(fn, 'rb') as f:
obj = load(f)
pprint.pprint(obj)


if __name__ == "__main__":
_main()
61 changes: 58 additions & 3 deletions Lib/test/test_pickle.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -699,6 +702,58 @@ def test_multiprocessing_exceptions(self):
('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
Expand Down
Loading