Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 1 | import pickle |
Tim Peters | 8ecfc8e | 2003-01-27 18:51:48 +0000 | [diff] [blame] | 2 | import pickletools |
Benjamin Peterson | ee8712c | 2008-05-20 21:35:26 +0000 | [diff] [blame] | 3 | from test import support |
Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 4 | from test.pickletester import AbstractPickleTests |
Martin Panter | 19e69c5 | 2015-11-14 12:46:42 +0000 | [diff] [blame] | 5 | import unittest |
Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 6 | |
Serhiy Storchaka | 6545256 | 2017-11-15 14:01:08 +0200 | [diff] [blame] | 7 | class OptimizedPickleTests(AbstractPickleTests): |
Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 8 | |
Antoine Pitrou | 91f4380 | 2019-05-26 17:10:09 +0200 | [diff] [blame] | 9 | def dumps(self, arg, proto=None, **kwargs): |
| 10 | return pickletools.optimize(pickle.dumps(arg, proto, **kwargs)) |
Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 11 | |
Alexander Belopolsky | 1ce92dc | 2011-02-24 19:40:09 +0000 | [diff] [blame] | 12 | def loads(self, buf, **kwds): |
| 13 | return pickle.loads(buf, **kwds) |
Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 14 | |
Antoine Pitrou | d9dfaa9 | 2009-06-04 20:32:06 +0000 | [diff] [blame] | 15 | # Test relies on precise output of dumps() |
| 16 | test_pickle_to_2x = None |
| 17 | |
Olivier Grisel | 3cd7c6e | 2018-01-06 16:18:54 +0100 | [diff] [blame] | 18 | # Test relies on writing by chunks into a file object. |
| 19 | test_framed_write_sizes_with_delayed_writer = None |
| 20 | |
Serhiy Storchaka | 05dadcf | 2014-12-16 18:00:56 +0200 | [diff] [blame] | 21 | def test_optimize_long_binget(self): |
| 22 | data = [str(i) for i in range(257)] |
| 23 | data.append(data[-1]) |
| 24 | for proto in range(pickle.HIGHEST_PROTOCOL + 1): |
| 25 | pickled = pickle.dumps(data, proto) |
| 26 | unpickled = pickle.loads(pickled) |
| 27 | self.assertEqual(unpickled, data) |
| 28 | self.assertIs(unpickled[-1], unpickled[-2]) |
| 29 | |
| 30 | pickled2 = pickletools.optimize(pickled) |
| 31 | unpickled2 = pickle.loads(pickled2) |
| 32 | self.assertEqual(unpickled2, data) |
| 33 | self.assertIs(unpickled2[-1], unpickled2[-2]) |
| 34 | self.assertNotIn(pickle.LONG_BINGET, pickled2) |
| 35 | self.assertNotIn(pickle.LONG_BINPUT, pickled2) |
| 36 | |
| 37 | def test_optimize_binput_and_memoize(self): |
| 38 | pickled = (b'\x80\x04\x95\x15\x00\x00\x00\x00\x00\x00\x00' |
| 39 | b']\x94(\x8c\x04spamq\x01\x8c\x03ham\x94h\x02e.') |
| 40 | # 0: \x80 PROTO 4 |
| 41 | # 2: \x95 FRAME 21 |
| 42 | # 11: ] EMPTY_LIST |
| 43 | # 12: \x94 MEMOIZE |
| 44 | # 13: ( MARK |
| 45 | # 14: \x8c SHORT_BINUNICODE 'spam' |
| 46 | # 20: q BINPUT 1 |
| 47 | # 22: \x8c SHORT_BINUNICODE 'ham' |
| 48 | # 27: \x94 MEMOIZE |
| 49 | # 28: h BINGET 2 |
| 50 | # 30: e APPENDS (MARK at 13) |
| 51 | # 31: . STOP |
| 52 | self.assertIn(pickle.BINPUT, pickled) |
| 53 | unpickled = pickle.loads(pickled) |
| 54 | self.assertEqual(unpickled, ['spam', 'ham', 'ham']) |
| 55 | self.assertIs(unpickled[1], unpickled[2]) |
| 56 | |
| 57 | pickled2 = pickletools.optimize(pickled) |
| 58 | unpickled2 = pickle.loads(pickled2) |
| 59 | self.assertEqual(unpickled2, ['spam', 'ham', 'ham']) |
| 60 | self.assertIs(unpickled2[1], unpickled2[2]) |
| 61 | self.assertNotIn(pickle.BINPUT, pickled2) |
| 62 | |
Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 63 | |
Martin Panter | 19e69c5 | 2015-11-14 12:46:42 +0000 | [diff] [blame] | 64 | class MiscTestCase(unittest.TestCase): |
| 65 | def test__all__(self): |
Victor Stinner | fbf43f0 | 2020-08-17 07:20:40 +0200 | [diff] [blame] | 66 | not_exported = { |
| 67 | 'bytes_types', |
| 68 | 'UP_TO_NEWLINE', 'TAKEN_FROM_ARGUMENT1', |
| 69 | 'TAKEN_FROM_ARGUMENT4', 'TAKEN_FROM_ARGUMENT4U', |
| 70 | 'TAKEN_FROM_ARGUMENT8U', 'ArgumentDescriptor', |
| 71 | 'read_uint1', 'read_uint2', 'read_int4', 'read_uint4', |
| 72 | 'read_uint8', 'read_stringnl', 'read_stringnl_noescape', |
| 73 | 'read_stringnl_noescape_pair', 'read_string1', |
| 74 | 'read_string4', 'read_bytes1', 'read_bytes4', |
| 75 | 'read_bytes8', 'read_bytearray8', 'read_unicodestringnl', |
| 76 | 'read_unicodestring1', 'read_unicodestring4', |
| 77 | 'read_unicodestring8', 'read_decimalnl_short', |
| 78 | 'read_decimalnl_long', 'read_floatnl', 'read_float8', |
| 79 | 'read_long1', 'read_long4', |
| 80 | 'uint1', 'uint2', 'int4', 'uint4', 'uint8', 'stringnl', |
| 81 | 'stringnl_noescape', 'stringnl_noescape_pair', 'string1', |
| 82 | 'string4', 'bytes1', 'bytes4', 'bytes8', 'bytearray8', |
| 83 | 'unicodestringnl', 'unicodestring1', 'unicodestring4', |
| 84 | 'unicodestring8', 'decimalnl_short', 'decimalnl_long', |
| 85 | 'floatnl', 'float8', 'long1', 'long4', |
| 86 | 'StackObject', |
| 87 | 'pyint', 'pylong', 'pyinteger_or_bool', 'pybool', 'pyfloat', |
| 88 | 'pybytes_or_str', 'pystring', 'pybytes', 'pybytearray', |
| 89 | 'pyunicode', 'pynone', 'pytuple', 'pylist', 'pydict', |
| 90 | 'pyset', 'pyfrozenset', 'pybuffer', 'anyobject', |
| 91 | 'markobject', 'stackslice', 'OpcodeInfo', 'opcodes', |
| 92 | 'code2op', |
| 93 | } |
| 94 | support.check__all__(self, pickletools, not_exported=not_exported) |
Martin Panter | 19e69c5 | 2015-11-14 12:46:42 +0000 | [diff] [blame] | 95 | |
| 96 | |
Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 97 | def test_main(): |
Benjamin Peterson | ee8712c | 2008-05-20 21:35:26 +0000 | [diff] [blame] | 98 | support.run_unittest(OptimizedPickleTests) |
Martin Panter | 19e69c5 | 2015-11-14 12:46:42 +0000 | [diff] [blame] | 99 | support.run_unittest(MiscTestCase) |
Benjamin Peterson | ee8712c | 2008-05-20 21:35:26 +0000 | [diff] [blame] | 100 | support.run_doctest(pickletools) |
Christian Heimes | 126d29a | 2008-02-11 22:57:17 +0000 | [diff] [blame] | 101 | |
| 102 | |
| 103 | if __name__ == "__main__": |
| 104 | test_main() |