Jeremy Hylton | be467e5 | 2000-09-15 15:14:51 +0000 | [diff] [blame] | 1 | import pickle |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 2 | import io |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 3 | import collections |
Serhiy Storchaka | 5bbd231 | 2014-12-16 19:39:08 +0200 | [diff] [blame] | 4 | import struct |
| 5 | import sys |
Tim Peters | 47a6b13 | 2003-01-28 22:34:11 +0000 | [diff] [blame] | 6 | |
Serhiy Storchaka | 5bbd231 | 2014-12-16 19:39:08 +0200 | [diff] [blame] | 7 | import unittest |
Benjamin Peterson | ee8712c | 2008-05-20 21:35:26 +0000 | [diff] [blame] | 8 | from test import support |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 9 | |
Tim Peters | 47a6b13 | 2003-01-28 22:34:11 +0000 | [diff] [blame] | 10 | from test.pickletester import AbstractPickleTests |
Tim Peters | 47a6b13 | 2003-01-28 22:34:11 +0000 | [diff] [blame] | 11 | from test.pickletester import AbstractPickleModuleTests |
| 12 | from test.pickletester import AbstractPersistentPicklerTests |
Collin Winter | 771d834 | 2009-04-16 03:18:06 +0000 | [diff] [blame] | 13 | from test.pickletester import AbstractPicklerUnpicklerObjectTests |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 14 | from test.pickletester import AbstractDispatchTableTests |
Antoine Pitrou | 82be19f | 2011-08-29 23:09:33 +0200 | [diff] [blame] | 15 | from test.pickletester import BigmemPickleTests |
Tim Peters | 47a6b13 | 2003-01-28 22:34:11 +0000 | [diff] [blame] | 16 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 17 | try: |
| 18 | import _pickle |
| 19 | has_c_implementation = True |
| 20 | except ImportError: |
| 21 | has_c_implementation = False |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 22 | |
Guido van Rossum | 98297ee | 2007-11-06 21:34:58 +0000 | [diff] [blame] | 23 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 24 | class PickleTests(AbstractPickleModuleTests): |
| 25 | pass |
Guido van Rossum | 5d9113d | 2003-01-29 17:58:45 +0000 | [diff] [blame] | 26 | |
Tim Peters | e0c446b | 2001-10-18 21:57:37 +0000 | [diff] [blame] | 27 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 28 | class PyPicklerTests(AbstractPickleTests): |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 29 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 30 | pickler = pickle._Pickler |
| 31 | unpickler = pickle._Unpickler |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 32 | |
Guido van Rossum | f416981 | 2008-03-17 22:56:06 +0000 | [diff] [blame] | 33 | def dumps(self, arg, proto=None): |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 34 | f = io.BytesIO() |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 35 | p = self.pickler(f, proto) |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 36 | p.dump(arg) |
| 37 | f.seek(0) |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 38 | return bytes(f.read()) |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 39 | |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 40 | def loads(self, buf, **kwds): |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 41 | f = io.BytesIO(buf) |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 42 | u = self.unpickler(f, **kwds) |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 43 | return u.load() |
| 44 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 45 | |
Antoine Pitrou | 82be19f | 2011-08-29 23:09:33 +0200 | [diff] [blame] | 46 | class InMemoryPickleTests(AbstractPickleTests, BigmemPickleTests): |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 47 | |
| 48 | pickler = pickle._Pickler |
| 49 | unpickler = pickle._Unpickler |
| 50 | |
Antoine Pitrou | 82be19f | 2011-08-29 23:09:33 +0200 | [diff] [blame] | 51 | def dumps(self, arg, protocol=None): |
| 52 | return pickle.dumps(arg, protocol) |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 53 | |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 54 | def loads(self, buf, **kwds): |
| 55 | return pickle.loads(buf, **kwds) |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 56 | |
| 57 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 58 | class PyPersPicklerTests(AbstractPersistentPicklerTests): |
| 59 | |
| 60 | pickler = pickle._Pickler |
| 61 | unpickler = pickle._Unpickler |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 62 | |
Guido van Rossum | f416981 | 2008-03-17 22:56:06 +0000 | [diff] [blame] | 63 | def dumps(self, arg, proto=None): |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 64 | class PersPickler(self.pickler): |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 65 | def persistent_id(subself, obj): |
| 66 | return self.persistent_id(obj) |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 67 | f = io.BytesIO() |
Guido van Rossum | 9d32bb1 | 2003-01-28 03:51:53 +0000 | [diff] [blame] | 68 | p = PersPickler(f, proto) |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 69 | p.dump(arg) |
| 70 | f.seek(0) |
| 71 | return f.read() |
| 72 | |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 73 | def loads(self, buf, **kwds): |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 74 | class PersUnpickler(self.unpickler): |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 75 | def persistent_load(subself, obj): |
| 76 | return self.persistent_load(obj) |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 77 | f = io.BytesIO(buf) |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 78 | u = PersUnpickler(f, **kwds) |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 79 | return u.load() |
| 80 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 81 | |
Collin Winter | 771d834 | 2009-04-16 03:18:06 +0000 | [diff] [blame] | 82 | class PyPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests): |
| 83 | |
| 84 | pickler_class = pickle._Pickler |
| 85 | unpickler_class = pickle._Unpickler |
| 86 | |
| 87 | |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 88 | class PyDispatchTableTests(AbstractDispatchTableTests): |
Alexandre Vassalotti | d05c9ff | 2013-12-07 01:09:27 -0800 | [diff] [blame] | 89 | |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 90 | pickler_class = pickle._Pickler |
Alexandre Vassalotti | d05c9ff | 2013-12-07 01:09:27 -0800 | [diff] [blame] | 91 | |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 92 | def get_dispatch_table(self): |
| 93 | return pickle.dispatch_table.copy() |
| 94 | |
| 95 | |
| 96 | class PyChainDispatchTableTests(AbstractDispatchTableTests): |
Alexandre Vassalotti | d05c9ff | 2013-12-07 01:09:27 -0800 | [diff] [blame] | 97 | |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 98 | pickler_class = pickle._Pickler |
Alexandre Vassalotti | d05c9ff | 2013-12-07 01:09:27 -0800 | [diff] [blame] | 99 | |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 100 | def get_dispatch_table(self): |
| 101 | return collections.ChainMap({}, pickle.dispatch_table) |
| 102 | |
| 103 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 104 | if has_c_implementation: |
| 105 | class CPicklerTests(PyPicklerTests): |
| 106 | pickler = _pickle.Pickler |
| 107 | unpickler = _pickle.Unpickler |
| 108 | |
| 109 | class CPersPicklerTests(PyPersPicklerTests): |
| 110 | pickler = _pickle.Pickler |
| 111 | unpickler = _pickle.Unpickler |
| 112 | |
Collin Winter | 771d834 | 2009-04-16 03:18:06 +0000 | [diff] [blame] | 113 | class CDumpPickle_LoadPickle(PyPicklerTests): |
| 114 | pickler = _pickle.Pickler |
| 115 | unpickler = pickle._Unpickler |
| 116 | |
| 117 | class DumpPickle_CLoadPickle(PyPicklerTests): |
| 118 | pickler = pickle._Pickler |
| 119 | unpickler = _pickle.Unpickler |
| 120 | |
| 121 | class CPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests): |
| 122 | pickler_class = _pickle.Pickler |
| 123 | unpickler_class = _pickle.Unpickler |
| 124 | |
Christian Heimes | a24b4d2 | 2013-07-01 15:17:45 +0200 | [diff] [blame] | 125 | def test_issue18339(self): |
| 126 | unpickler = self.unpickler_class(io.BytesIO()) |
Christian Heimes | 2178248 | 2013-07-01 23:00:13 +0200 | [diff] [blame] | 127 | with self.assertRaises(TypeError): |
| 128 | unpickler.memo = object |
Christian Heimes | a24b4d2 | 2013-07-01 15:17:45 +0200 | [diff] [blame] | 129 | # used to cause a segfault |
Christian Heimes | 2178248 | 2013-07-01 23:00:13 +0200 | [diff] [blame] | 130 | with self.assertRaises(ValueError): |
| 131 | unpickler.memo = {-1: None} |
Christian Heimes | a24b4d2 | 2013-07-01 15:17:45 +0200 | [diff] [blame] | 132 | unpickler.memo = {1: None} |
| 133 | |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 134 | class CDispatchTableTests(AbstractDispatchTableTests): |
| 135 | pickler_class = pickle.Pickler |
| 136 | def get_dispatch_table(self): |
| 137 | return pickle.dispatch_table.copy() |
| 138 | |
| 139 | class CChainDispatchTableTests(AbstractDispatchTableTests): |
| 140 | pickler_class = pickle.Pickler |
| 141 | def get_dispatch_table(self): |
| 142 | return collections.ChainMap({}, pickle.dispatch_table) |
| 143 | |
Serhiy Storchaka | 5bbd231 | 2014-12-16 19:39:08 +0200 | [diff] [blame] | 144 | @support.cpython_only |
| 145 | class SizeofTests(unittest.TestCase): |
| 146 | check_sizeof = support.check_sizeof |
| 147 | |
| 148 | def test_pickler(self): |
| 149 | basesize = support.calcobjsize('5P2n3i2n3iP') |
| 150 | p = _pickle.Pickler(io.BytesIO()) |
| 151 | self.assertEqual(object.__sizeof__(p), basesize) |
| 152 | MT_size = struct.calcsize('3nP0n') |
| 153 | ME_size = struct.calcsize('Pn0P') |
| 154 | check = self.check_sizeof |
| 155 | check(p, basesize + |
| 156 | MT_size + 8 * ME_size + # Minimal memo table size. |
| 157 | sys.getsizeof(b'x'*4096)) # Minimal write buffer size. |
| 158 | for i in range(6): |
| 159 | p.dump(chr(i)) |
| 160 | check(p, basesize + |
| 161 | MT_size + 32 * ME_size + # Size of memo table required to |
| 162 | # save references to 6 objects. |
| 163 | 0) # Write buffer is cleared after every dump(). |
| 164 | |
| 165 | def test_unpickler(self): |
| 166 | basesize = support.calcobjsize('2Pn2P 2P2n2i5P 2P3n6P2n2i') |
| 167 | unpickler = _pickle.Unpickler |
| 168 | P = struct.calcsize('P') # Size of memo table entry. |
| 169 | n = struct.calcsize('n') # Size of mark table entry. |
| 170 | check = self.check_sizeof |
| 171 | for encoding in 'ASCII', 'UTF-16', 'latin-1': |
| 172 | for errors in 'strict', 'replace': |
| 173 | u = unpickler(io.BytesIO(), |
| 174 | encoding=encoding, errors=errors) |
| 175 | self.assertEqual(object.__sizeof__(u), basesize) |
| 176 | check(u, basesize + |
| 177 | 32 * P + # Minimal memo table size. |
| 178 | len(encoding) + 1 + len(errors) + 1) |
| 179 | |
| 180 | stdsize = basesize + len('ASCII') + 1 + len('strict') + 1 |
| 181 | def check_unpickler(data, memo_size, marks_size): |
| 182 | dump = pickle.dumps(data) |
| 183 | u = unpickler(io.BytesIO(dump), |
| 184 | encoding='ASCII', errors='strict') |
| 185 | u.load() |
| 186 | check(u, stdsize + memo_size * P + marks_size * n) |
| 187 | |
| 188 | check_unpickler(0, 32, 0) |
| 189 | # 20 is minimal non-empty mark stack size. |
| 190 | check_unpickler([0] * 100, 32, 20) |
| 191 | # 128 is memo table size required to save references to 100 objects. |
| 192 | check_unpickler([chr(i) for i in range(100)], 128, 20) |
| 193 | def recurse(deep): |
| 194 | data = 0 |
| 195 | for i in range(deep): |
| 196 | data = [data, data] |
| 197 | return data |
| 198 | check_unpickler(recurse(0), 32, 0) |
| 199 | check_unpickler(recurse(1), 32, 20) |
| 200 | check_unpickler(recurse(20), 32, 58) |
| 201 | check_unpickler(recurse(50), 64, 58) |
| 202 | check_unpickler(recurse(100), 128, 134) |
| 203 | |
| 204 | u = unpickler(io.BytesIO(pickle.dumps('a', 0)), |
| 205 | encoding='ASCII', errors='strict') |
| 206 | u.load() |
| 207 | check(u, stdsize + 32 * P + 2 + 1) |
| 208 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 209 | |
Fred Drake | 694ed09 | 2001-12-19 16:42:15 +0000 | [diff] [blame] | 210 | def test_main(): |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 211 | tests = [PickleTests, PyPicklerTests, PyPersPicklerTests, |
| 212 | PyDispatchTableTests, PyChainDispatchTableTests] |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 213 | if has_c_implementation: |
Collin Winter | 771d834 | 2009-04-16 03:18:06 +0000 | [diff] [blame] | 214 | tests.extend([CPicklerTests, CPersPicklerTests, |
| 215 | CDumpPickle_LoadPickle, DumpPickle_CLoadPickle, |
| 216 | PyPicklerUnpicklerObjectTests, |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 217 | CPicklerUnpicklerObjectTests, |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 218 | CDispatchTableTests, CChainDispatchTableTests, |
Serhiy Storchaka | 5bbd231 | 2014-12-16 19:39:08 +0200 | [diff] [blame] | 219 | InMemoryPickleTests, SizeofTests]) |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 220 | support.run_unittest(*tests) |
Benjamin Peterson | ee8712c | 2008-05-20 21:35:26 +0000 | [diff] [blame] | 221 | support.run_doctest(pickle) |
Fred Drake | 694ed09 | 2001-12-19 16:42:15 +0000 | [diff] [blame] | 222 | |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 223 | if __name__ == "__main__": |
Fred Drake | 694ed09 | 2001-12-19 16:42:15 +0000 | [diff] [blame] | 224 | test_main() |