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