1import pickle
2import pickletools
3from test import support
4from test.pickletester import AbstractPickleTests
5from test.pickletester import AbstractPickleModuleTests
6import unittest
7
8class OptimizedPickleTests(AbstractPickleTests, AbstractPickleModuleTests):
9
10    def dumps(self, arg, proto=None):
11        return pickletools.optimize(pickle.dumps(arg, proto))
12
13    def loads(self, buf, **kwds):
14        return pickle.loads(buf, **kwds)
15
16    # Test relies on precise output of dumps()
17    test_pickle_to_2x = None
18
19    def test_optimize_long_binget(self):
20        data = [str(i) for i in range(257)]
21        data.append(data[-1])
22        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
23            pickled = pickle.dumps(data, proto)
24            unpickled = pickle.loads(pickled)
25            self.assertEqual(unpickled, data)
26            self.assertIs(unpickled[-1], unpickled[-2])
27
28            pickled2 = pickletools.optimize(pickled)
29            unpickled2 = pickle.loads(pickled2)
30            self.assertEqual(unpickled2, data)
31            self.assertIs(unpickled2[-1], unpickled2[-2])
32            self.assertNotIn(pickle.LONG_BINGET, pickled2)
33            self.assertNotIn(pickle.LONG_BINPUT, pickled2)
34
35    def test_optimize_binput_and_memoize(self):
36        pickled = (b'\x80\x04\x95\x15\x00\x00\x00\x00\x00\x00\x00'
37                   b']\x94(\x8c\x04spamq\x01\x8c\x03ham\x94h\x02e.')
38        #    0: \x80 PROTO      4
39        #    2: \x95 FRAME      21
40        #   11: ]    EMPTY_LIST
41        #   12: \x94 MEMOIZE
42        #   13: (    MARK
43        #   14: \x8c     SHORT_BINUNICODE 'spam'
44        #   20: q        BINPUT     1
45        #   22: \x8c     SHORT_BINUNICODE 'ham'
46        #   27: \x94     MEMOIZE
47        #   28: h        BINGET     2
48        #   30: e        APPENDS    (MARK at 13)
49        #   31: .    STOP
50        self.assertIn(pickle.BINPUT, pickled)
51        unpickled = pickle.loads(pickled)
52        self.assertEqual(unpickled, ['spam', 'ham', 'ham'])
53        self.assertIs(unpickled[1], unpickled[2])
54
55        pickled2 = pickletools.optimize(pickled)
56        unpickled2 = pickle.loads(pickled2)
57        self.assertEqual(unpickled2, ['spam', 'ham', 'ham'])
58        self.assertIs(unpickled2[1], unpickled2[2])
59        self.assertNotIn(pickle.BINPUT, pickled2)
60
61
62class MiscTestCase(unittest.TestCase):
63    def test__all__(self):
64        blacklist = {'bytes_types',
65                     'UP_TO_NEWLINE', 'TAKEN_FROM_ARGUMENT1',
66                     'TAKEN_FROM_ARGUMENT4', 'TAKEN_FROM_ARGUMENT4U',
67                     'TAKEN_FROM_ARGUMENT8U', 'ArgumentDescriptor',
68                     'read_uint1', 'read_uint2', 'read_int4', 'read_uint4',
69                     'read_uint8', 'read_stringnl', 'read_stringnl_noescape',
70                     'read_stringnl_noescape_pair', 'read_string1',
71                     'read_string4', 'read_bytes1', 'read_bytes4',
72                     'read_bytes8', 'read_unicodestringnl',
73                     'read_unicodestring1', 'read_unicodestring4',
74                     'read_unicodestring8', 'read_decimalnl_short',
75                     'read_decimalnl_long', 'read_floatnl', 'read_float8',
76                     'read_long1', 'read_long4',
77                     'uint1', 'uint2', 'int4', 'uint4', 'uint8', 'stringnl',
78                     'stringnl_noescape', 'stringnl_noescape_pair', 'string1',
79                     'string4', 'bytes1', 'bytes4', 'bytes8',
80                     'unicodestringnl', 'unicodestring1', 'unicodestring4',
81                     'unicodestring8', 'decimalnl_short', 'decimalnl_long',
82                     'floatnl', 'float8', 'long1', 'long4',
83                     'StackObject',
84                     'pyint', 'pylong', 'pyinteger_or_bool', 'pybool', 'pyfloat',
85                     'pybytes_or_str', 'pystring', 'pybytes', 'pyunicode',
86                     'pynone', 'pytuple', 'pylist', 'pydict', 'pyset',
87                     'pyfrozenset', 'anyobject', 'markobject', 'stackslice',
88                     'OpcodeInfo', 'opcodes', 'code2op',
89                     }
90        support.check__all__(self, pickletools, blacklist=blacklist)
91
92
93def test_main():
94    support.run_unittest(OptimizedPickleTests)
95    support.run_unittest(MiscTestCase)
96    support.run_doctest(pickletools)
97
98
99if __name__ == "__main__":
100    test_main()
101