blob: d54119470362ed46208087d38e56de42e7ab0c15 [file] [log] [blame]
Jeremy Hylton66426532001-10-15 21:38:56 +00001import unittest
Tim Peters4190fb82003-02-02 16:09:05 +00002import pickle
Tim Peters8587b3c2003-02-13 15:44:41 +00003import cPickle
Tim Peters31f119e2003-02-03 16:20:13 +00004import pickletools
Tim Peters3e667d52003-02-04 21:47:44 +00005import copy_reg
Tim Peters4190fb82003-02-02 16:09:05 +00006
Guido van Rossum3b0a3292002-08-09 16:38:32 +00007from test.test_support import TestFailed, have_unicode, TESTFN
Tim Peterse089c682001-04-10 03:41:41 +00008
Tim Petersee1a53c2003-02-02 02:57:53 +00009# Tests that try a number of pickle protocols should have a
10# for proto in protocols:
Tim Peters8587b3c2003-02-13 15:44:41 +000011# kind of outer loop.
12assert pickle.HIGHEST_PROTOCOL == cPickle.HIGHEST_PROTOCOL == 2
13protocols = range(pickle.HIGHEST_PROTOCOL + 1)
Tim Petersee1a53c2003-02-02 02:57:53 +000014
Tim Peters22e71712003-02-03 22:27:38 +000015
16# Return True if opcode code appears in the pickle, else False.
17def opcode_in_pickle(code, pickle):
18 for op, dummy, dummy in pickletools.genops(pickle):
19 if op.code == code:
20 return True
21 return False
22
Tim Peters8d2613a2003-02-11 16:40:16 +000023# Return the number of times opcode code appears in pickle.
24def count_opcode(code, pickle):
25 n = 0
26 for op, dummy, dummy in pickletools.genops(pickle):
27 if op.code == code:
28 n += 1
29 return n
30
Tim Peters3e667d52003-02-04 21:47:44 +000031# We can't very well test the extension registry without putting known stuff
32# in it, but we have to be careful to restore its original state. Code
33# should do this:
34#
35# e = ExtensionSaver(extension_code)
36# try:
37# fiddle w/ the extension registry's stuff for extension_code
38# finally:
39# e.restore()
40
41class ExtensionSaver:
42 # Remember current registration for code (if any), and remove it (if
43 # there is one).
44 def __init__(self, code):
45 self.code = code
46 if code in copy_reg._inverted_registry:
47 self.pair = copy_reg._inverted_registry[code]
48 copy_reg.remove_extension(self.pair[0], self.pair[1], code)
49 else:
50 self.pair = None
51
52 # Restore previous registration for code.
53 def restore(self):
54 code = self.code
55 curpair = copy_reg._inverted_registry.get(code)
56 if curpair is not None:
57 copy_reg.remove_extension(curpair[0], curpair[1], code)
58 pair = self.pair
59 if pair is not None:
60 copy_reg.add_extension(pair[0], pair[1], code)
61
Jeremy Hylton66426532001-10-15 21:38:56 +000062class C:
63 def __cmp__(self, other):
64 return cmp(self.__dict__, other.__dict__)
65
66import __main__
67__main__.C = C
68C.__module__ = "__main__"
69
70class myint(int):
71 def __init__(self, x):
72 self.str = str(x)
73
74class initarg(C):
Guido van Rossum1444f672001-12-19 16:38:29 +000075
Jeremy Hylton66426532001-10-15 21:38:56 +000076 def __init__(self, a, b):
77 self.a = a
78 self.b = b
79
80 def __getinitargs__(self):
81 return self.a, self.b
82
Guido van Rossum04a86612001-12-19 16:58:54 +000083class metaclass(type):
84 pass
85
86class use_metaclass(object):
87 __metaclass__ = metaclass
88
Tim Peters70b02d72003-02-02 17:26:40 +000089# DATA0 .. DATA2 are the pickles we expect under the various protocols, for
90# the object returned by create_data().
91# XXX DATA2 doesn't exist yet, as it's not fully implemented in cPickle.
Tim Petersee1a53c2003-02-02 02:57:53 +000092
Jeremy Hylton66426532001-10-15 21:38:56 +000093# break into multiple strings to avoid confusing font-lock-mode
Tim Peters70b02d72003-02-02 17:26:40 +000094DATA0 = """(lp1
Tim Peterse9358162001-01-22 22:05:20 +000095I0
96aL1L
Tim Peters461922a2001-04-09 20:07:05 +000097aF2
Tim Peterse9358162001-01-22 22:05:20 +000098ac__builtin__
99complex
Tim Peters461922a2001-04-09 20:07:05 +0000100p2
101""" + \
102"""(F3
103F0
104tRp3
105aI1
106aI-1
107aI255
108aI-255
109aI-256
110aI65535
111aI-65535
112aI-65536
113aI2147483647
114aI-2147483647
115aI-2147483648
116a""" + \
117"""(S'abc'
Tim Peterse9358162001-01-22 22:05:20 +0000118p4
119g4
Tim Peters461922a2001-04-09 20:07:05 +0000120""" + \
Guido van Rossum42f92da2001-04-16 00:28:21 +0000121"""(i__main__
Tim Peterse9358162001-01-22 22:05:20 +0000122C
123p5
Tim Peters461922a2001-04-09 20:07:05 +0000124""" + \
Tim Peterse9358162001-01-22 22:05:20 +0000125"""(dp6
126S'foo'
127p7
128I1
129sS'bar'
130p8
131I2
132sbg5
133tp9
134ag9
135aI5
136a.
137"""
138
Tim Peters70b02d72003-02-02 17:26:40 +0000139# Disassembly of DATA0.
140DATA0_DIS = """\
141 0: ( MARK
142 1: l LIST (MARK at 0)
143 2: p PUT 1
144 5: I INT 0
145 8: a APPEND
146 9: L LONG 1L
147 13: a APPEND
148 14: F FLOAT 2.0
149 17: a APPEND
150 18: c GLOBAL '__builtin__ complex'
151 39: p PUT 2
152 42: ( MARK
153 43: F FLOAT 3.0
154 46: F FLOAT 0.0
155 49: t TUPLE (MARK at 42)
156 50: R REDUCE
157 51: p PUT 3
158 54: a APPEND
159 55: I INT 1
160 58: a APPEND
161 59: I INT -1
162 63: a APPEND
163 64: I INT 255
164 69: a APPEND
165 70: I INT -255
166 76: a APPEND
167 77: I INT -256
168 83: a APPEND
169 84: I INT 65535
170 91: a APPEND
171 92: I INT -65535
172 100: a APPEND
173 101: I INT -65536
174 109: a APPEND
175 110: I INT 2147483647
176 122: a APPEND
177 123: I INT -2147483647
178 136: a APPEND
179 137: I INT -2147483648
180 150: a APPEND
181 151: ( MARK
182 152: S STRING 'abc'
183 159: p PUT 4
184 162: g GET 4
185 165: ( MARK
186 166: i INST '__main__ C' (MARK at 165)
187 178: p PUT 5
188 181: ( MARK
189 182: d DICT (MARK at 181)
190 183: p PUT 6
191 186: S STRING 'foo'
192 193: p PUT 7
193 196: I INT 1
194 199: s SETITEM
195 200: S STRING 'bar'
196 207: p PUT 8
197 210: I INT 2
198 213: s SETITEM
199 214: b BUILD
200 215: g GET 5
201 218: t TUPLE (MARK at 151)
202 219: p PUT 9
203 222: a APPEND
204 223: g GET 9
205 226: a APPEND
206 227: I INT 5
207 230: a APPEND
208 231: . STOP
209highest protocol among opcodes = 0
210"""
211
212DATA1 = (']q\x01(K\x00L1L\nG@\x00\x00\x00\x00\x00\x00\x00'
213 'c__builtin__\ncomplex\nq\x02(G@\x08\x00\x00\x00\x00\x00'
214 '\x00G\x00\x00\x00\x00\x00\x00\x00\x00tRq\x03K\x01J\xff\xff'
215 '\xff\xffK\xffJ\x01\xff\xff\xffJ\x00\xff\xff\xffM\xff\xff'
216 'J\x01\x00\xff\xffJ\x00\x00\xff\xffJ\xff\xff\xff\x7fJ\x01\x00'
217 '\x00\x80J\x00\x00\x00\x80(U\x03abcq\x04h\x04(c__main__\n'
218 'C\nq\x05oq\x06}q\x07(U\x03fooq\x08K\x01U\x03barq\tK\x02ubh'
219 '\x06tq\nh\nK\x05e.'
220 )
221
222# Disassembly of DATA1.
223DATA1_DIS = """\
224 0: ] EMPTY_LIST
225 1: q BINPUT 1
226 3: ( MARK
227 4: K BININT1 0
228 6: L LONG 1L
229 10: G BINFLOAT 2.0
230 19: c GLOBAL '__builtin__ complex'
231 40: q BINPUT 2
232 42: ( MARK
233 43: G BINFLOAT 3.0
234 52: G BINFLOAT 0.0
235 61: t TUPLE (MARK at 42)
236 62: R REDUCE
237 63: q BINPUT 3
238 65: K BININT1 1
239 67: J BININT -1
240 72: K BININT1 255
241 74: J BININT -255
242 79: J BININT -256
243 84: M BININT2 65535
244 87: J BININT -65535
245 92: J BININT -65536
246 97: J BININT 2147483647
247 102: J BININT -2147483647
248 107: J BININT -2147483648
249 112: ( MARK
250 113: U SHORT_BINSTRING 'abc'
251 118: q BINPUT 4
252 120: h BINGET 4
253 122: ( MARK
254 123: c GLOBAL '__main__ C'
255 135: q BINPUT 5
256 137: o OBJ (MARK at 122)
257 138: q BINPUT 6
258 140: } EMPTY_DICT
259 141: q BINPUT 7
260 143: ( MARK
261 144: U SHORT_BINSTRING 'foo'
262 149: q BINPUT 8
263 151: K BININT1 1
264 153: U SHORT_BINSTRING 'bar'
265 158: q BINPUT 9
266 160: K BININT1 2
267 162: u SETITEMS (MARK at 143)
268 163: b BUILD
269 164: h BINGET 6
270 166: t TUPLE (MARK at 112)
271 167: q BINPUT 10
272 169: h BINGET 10
273 171: K BININT1 5
274 173: e APPENDS (MARK at 3)
275 174: . STOP
276highest protocol among opcodes = 1
277"""
Tim Peterse0c446b2001-10-18 21:57:37 +0000278
Jeremy Hylton66426532001-10-15 21:38:56 +0000279def create_data():
Tim Peterse9358162001-01-22 22:05:20 +0000280 c = C()
281 c.foo = 1
282 c.bar = 2
283 x = [0, 1L, 2.0, 3.0+0j]
Tim Peters461922a2001-04-09 20:07:05 +0000284 # Append some integer test cases at cPickle.c's internal size
285 # cutoffs.
286 uint1max = 0xff
287 uint2max = 0xffff
288 int4max = 0x7fffffff
289 x.extend([1, -1,
290 uint1max, -uint1max, -uint1max-1,
291 uint2max, -uint2max, -uint2max-1,
292 int4max, -int4max, -int4max-1])
Tim Peterse9358162001-01-22 22:05:20 +0000293 y = ('abc', 'abc', c, c)
294 x.append(y)
295 x.append(y)
296 x.append(5)
Jeremy Hylton66426532001-10-15 21:38:56 +0000297 return x
Tim Petersc58440f2001-04-09 17:16:31 +0000298
Jeremy Hylton66426532001-10-15 21:38:56 +0000299class AbstractPickleTests(unittest.TestCase):
Tim Peters70b02d72003-02-02 17:26:40 +0000300 # Subclass must define self.dumps, self.loads, self.error.
Tim Petersc58440f2001-04-09 17:16:31 +0000301
Jeremy Hylton66426532001-10-15 21:38:56 +0000302 _testdata = create_data()
Tim Petersc58440f2001-04-09 17:16:31 +0000303
Jeremy Hylton66426532001-10-15 21:38:56 +0000304 def setUp(self):
Tim Peterse9358162001-01-22 22:05:20 +0000305 pass
Tim Petersc58440f2001-04-09 17:16:31 +0000306
Jeremy Hylton66426532001-10-15 21:38:56 +0000307 def test_misc(self):
308 # test various datatypes not tested by testdata
Tim Peters70b02d72003-02-02 17:26:40 +0000309 for proto in protocols:
310 x = myint(4)
311 s = self.dumps(x, proto)
312 y = self.loads(s)
313 self.assertEqual(x, y)
Tim Peterse9358162001-01-22 22:05:20 +0000314
Tim Peters70b02d72003-02-02 17:26:40 +0000315 x = (1, ())
316 s = self.dumps(x, proto)
317 y = self.loads(s)
318 self.assertEqual(x, y)
Tim Peterse9358162001-01-22 22:05:20 +0000319
Tim Peters70b02d72003-02-02 17:26:40 +0000320 x = initarg(1, x)
321 s = self.dumps(x, proto)
322 y = self.loads(s)
323 self.assertEqual(x, y)
Tim Peterse9358162001-01-22 22:05:20 +0000324
Jeremy Hylton66426532001-10-15 21:38:56 +0000325 # XXX test __reduce__ protocol?
326
Tim Peters70b02d72003-02-02 17:26:40 +0000327 def test_roundtrip_equality(self):
328 expected = self._testdata
329 for proto in protocols:
330 s = self.dumps(expected, proto)
331 got = self.loads(s)
332 self.assertEqual(expected, got)
Jeremy Hylton66426532001-10-15 21:38:56 +0000333
Tim Peters70b02d72003-02-02 17:26:40 +0000334 def test_load_from_canned_string(self):
335 expected = self._testdata
336 for canned in DATA0, DATA1:
337 got = self.loads(canned)
338 self.assertEqual(expected, got)
Jeremy Hylton66426532001-10-15 21:38:56 +0000339
Tim Peters70b02d72003-02-02 17:26:40 +0000340 # There are gratuitous differences between pickles produced by
341 # pickle and cPickle, largely because cPickle starts PUT indices at
342 # 1 and pickle starts them at 0. See XXX comment in cPickle's put2() --
343 # there's a comment with an exclamation point there whose meaning
344 # is a mystery. cPickle also suppresses PUT for objects with a refcount
345 # of 1.
346 def dont_test_disassembly(self):
347 from cStringIO import StringIO
348 from pickletools import dis
349
350 for proto, expected in (0, DATA0_DIS), (1, DATA1_DIS):
351 s = self.dumps(self._testdata, proto)
352 filelike = StringIO()
353 dis(s, out=filelike)
354 got = filelike.getvalue()
355 self.assertEqual(expected, got)
Jeremy Hylton66426532001-10-15 21:38:56 +0000356
357 def test_recursive_list(self):
358 l = []
359 l.append(l)
Tim Peters70b02d72003-02-02 17:26:40 +0000360 for proto in protocols:
361 s = self.dumps(l, proto)
362 x = self.loads(s)
363 self.assertEqual(x, l)
364 self.assertEqual(x, x[0])
365 self.assertEqual(id(x), id(x[0]))
Jeremy Hylton66426532001-10-15 21:38:56 +0000366
367 def test_recursive_dict(self):
368 d = {}
369 d[1] = d
Tim Peters70b02d72003-02-02 17:26:40 +0000370 for proto in protocols:
371 s = self.dumps(d, proto)
372 x = self.loads(s)
373 self.assertEqual(x, d)
374 self.assertEqual(x[1], x)
375 self.assertEqual(id(x[1]), id(x))
Jeremy Hylton66426532001-10-15 21:38:56 +0000376
377 def test_recursive_inst(self):
378 i = C()
379 i.attr = i
Tim Peters70b02d72003-02-02 17:26:40 +0000380 for proto in protocols:
381 s = self.dumps(i, 2)
382 x = self.loads(s)
383 self.assertEqual(x, i)
384 self.assertEqual(x.attr, x)
385 self.assertEqual(id(x.attr), id(x))
Jeremy Hylton66426532001-10-15 21:38:56 +0000386
387 def test_recursive_multi(self):
388 l = []
389 d = {1:l}
390 i = C()
391 i.attr = d
392 l.append(i)
Tim Peters70b02d72003-02-02 17:26:40 +0000393 for proto in protocols:
394 s = self.dumps(l, proto)
395 x = self.loads(s)
396 self.assertEqual(x, l)
397 self.assertEqual(x[0], i)
398 self.assertEqual(x[0].attr, d)
399 self.assertEqual(x[0].attr[1], x)
400 self.assertEqual(x[0].attr[1][0], i)
401 self.assertEqual(x[0].attr[1][0].attr, d)
Jeremy Hylton66426532001-10-15 21:38:56 +0000402
403 def test_garyp(self):
404 self.assertRaises(self.error, self.loads, 'garyp')
405
406 def test_insecure_strings(self):
407 insecure = ["abc", "2 + 2", # not quoted
Martin v. Löwis8a8da792002-08-14 07:46:28 +0000408 #"'abc' + 'def'", # not a single quoted string
Jeremy Hylton66426532001-10-15 21:38:56 +0000409 "'abc", # quote is not closed
410 "'abc\"", # open quote and close quote don't match
411 "'abc' ?", # junk after close quote
Martin v. Löwiseb3f00a2002-08-14 08:22:50 +0000412 "'\\'", # trailing backslash
Jeremy Hylton66426532001-10-15 21:38:56 +0000413 # some tests of the quoting rules
Martin v. Löwis8a8da792002-08-14 07:46:28 +0000414 #"'abc\"\''",
415 #"'\\\\a\'\'\'\\\'\\\\\''",
Jeremy Hylton66426532001-10-15 21:38:56 +0000416 ]
417 for s in insecure:
418 buf = "S" + s + "\012p0\012."
419 self.assertRaises(ValueError, self.loads, buf)
420
Martin v. Löwis339d0f72001-08-17 18:39:25 +0000421 if have_unicode:
Jeremy Hylton66426532001-10-15 21:38:56 +0000422 def test_unicode(self):
423 endcases = [unicode(''), unicode('<\\u>'), unicode('<\\\u1234>'),
424 unicode('<\n>'), unicode('<\\>')]
Tim Petersee1a53c2003-02-02 02:57:53 +0000425 for proto in protocols:
426 for u in endcases:
427 p = self.dumps(u, proto)
428 u2 = self.loads(p)
429 self.assertEqual(u2, u)
Tim Peterse089c682001-04-10 03:41:41 +0000430
Jeremy Hylton66426532001-10-15 21:38:56 +0000431 def test_ints(self):
432 import sys
Tim Petersee1a53c2003-02-02 02:57:53 +0000433 for proto in protocols:
434 n = sys.maxint
435 while n:
436 for expected in (-n, n):
437 s = self.dumps(expected, proto)
438 n2 = self.loads(s)
439 self.assertEqual(expected, n2)
440 n = n >> 1
Tim Peters19ef62d2001-08-28 22:21:18 +0000441
Jeremy Hylton66426532001-10-15 21:38:56 +0000442 def test_maxint64(self):
443 maxint64 = (1L << 63) - 1
444 data = 'I' + str(maxint64) + '\n.'
445 got = self.loads(data)
446 self.assertEqual(got, maxint64)
447
448 # Try too with a bogus literal.
449 data = 'I' + str(maxint64) + 'JUNK\n.'
450 self.assertRaises(ValueError, self.loads, data)
451
Tim Petersee1a53c2003-02-02 02:57:53 +0000452 def test_long(self):
453 for proto in protocols:
Tim Petersbf2674b2003-02-02 07:51:32 +0000454 # 256 bytes is where LONG4 begins.
Tim Petersee1a53c2003-02-02 02:57:53 +0000455 for nbits in 1, 8, 8*254, 8*255, 8*256, 8*257:
456 nbase = 1L << nbits
457 for npos in nbase-1, nbase, nbase+1:
458 for n in npos, -npos:
459 pickle = self.dumps(n, proto)
460 got = self.loads(pickle)
461 self.assertEqual(n, got)
462 # Try a monster. This is quadratic-time in protos 0 & 1, so don't
463 # bother with those.
Tim Petersee1a53c2003-02-02 02:57:53 +0000464 nbase = long("deadbeeffeedface", 16)
465 nbase += nbase << 1000000
466 for n in nbase, -nbase:
Tim Petersee1a53c2003-02-02 02:57:53 +0000467 p = self.dumps(n, 2)
Tim Petersee1a53c2003-02-02 02:57:53 +0000468 got = self.loads(p)
Tim Petersee1a53c2003-02-02 02:57:53 +0000469 self.assertEqual(n, got)
470
Jeremy Hylton66426532001-10-15 21:38:56 +0000471 def test_reduce(self):
Tim Peters19ef62d2001-08-28 22:21:18 +0000472 pass
Jeremy Hylton66426532001-10-15 21:38:56 +0000473
474 def test_getinitargs(self):
475 pass
476
Guido van Rossum04a86612001-12-19 16:58:54 +0000477 def test_metaclass(self):
478 a = use_metaclass()
Tim Peters70b02d72003-02-02 17:26:40 +0000479 for proto in protocols:
480 s = self.dumps(a, proto)
481 b = self.loads(s)
482 self.assertEqual(a.__class__, b.__class__)
Guido van Rossum04a86612001-12-19 16:58:54 +0000483
Michael W. Hudson7bb466a2002-03-05 13:27:58 +0000484 def test_structseq(self):
485 import time
Michael W. Hudson0e025302002-03-06 17:11:18 +0000486 import os
Tim Peters70b02d72003-02-02 17:26:40 +0000487
488 t = time.localtime()
489 for proto in protocols:
490 s = self.dumps(t, proto)
Michael W. Hudson0e025302002-03-06 17:11:18 +0000491 u = self.loads(s)
492 self.assertEqual(t, u)
Tim Peters70b02d72003-02-02 17:26:40 +0000493 if hasattr(os, "stat"):
494 t = os.stat(os.curdir)
495 s = self.dumps(t, proto)
496 u = self.loads(s)
497 self.assertEqual(t, u)
498 if hasattr(os, "statvfs"):
499 t = os.statvfs(os.curdir)
500 s = self.dumps(t, proto)
501 u = self.loads(s)
502 self.assertEqual(t, u)
Michael W. Hudson7bb466a2002-03-05 13:27:58 +0000503
Guido van Rossumd6c9e632003-01-28 03:49:52 +0000504 # Tests for protocol 2
505
Tim Peters4190fb82003-02-02 16:09:05 +0000506 def test_proto(self):
507 build_none = pickle.NONE + pickle.STOP
508 for proto in protocols:
509 expected = build_none
510 if proto >= 2:
511 expected = pickle.PROTO + chr(proto) + expected
512 p = self.dumps(None, proto)
513 self.assertEqual(p, expected)
514
515 oob = protocols[-1] + 1 # a future protocol
516 badpickle = pickle.PROTO + chr(oob) + build_none
517 try:
518 self.loads(badpickle)
519 except ValueError, detail:
520 self.failUnless(str(detail).startswith(
521 "unsupported pickle protocol"))
522 else:
523 self.fail("expected bad protocol number to raise ValueError")
524
Guido van Rossumd6c9e632003-01-28 03:49:52 +0000525 def test_long1(self):
526 x = 12345678910111213141516178920L
Tim Peters61bf2572003-02-03 21:31:22 +0000527 for proto in protocols:
528 s = self.dumps(x, proto)
529 y = self.loads(s)
530 self.assertEqual(x, y)
Tim Peters22e71712003-02-03 22:27:38 +0000531 self.assertEqual(opcode_in_pickle(pickle.LONG1, s), proto >= 2)
Guido van Rossumd6c9e632003-01-28 03:49:52 +0000532
533 def test_long4(self):
534 x = 12345678910111213141516178920L << (256*8)
Tim Peters61bf2572003-02-03 21:31:22 +0000535 for proto in protocols:
536 s = self.dumps(x, proto)
537 y = self.loads(s)
538 self.assertEqual(x, y)
Tim Peters22e71712003-02-03 22:27:38 +0000539 self.assertEqual(opcode_in_pickle(pickle.LONG4, s), proto >= 2)
Guido van Rossumd6c9e632003-01-28 03:49:52 +0000540
Guido van Rossum44f0ea52003-01-28 04:14:51 +0000541 def test_short_tuples(self):
Tim Peters1d63c9f2003-02-02 20:29:39 +0000542 # Map (proto, len(tuple)) to expected opcode.
543 expected_opcode = {(0, 0): pickle.TUPLE,
544 (0, 1): pickle.TUPLE,
545 (0, 2): pickle.TUPLE,
546 (0, 3): pickle.TUPLE,
547 (0, 4): pickle.TUPLE,
548
549 (1, 0): pickle.EMPTY_TUPLE,
550 (1, 1): pickle.TUPLE,
551 (1, 2): pickle.TUPLE,
552 (1, 3): pickle.TUPLE,
553 (1, 4): pickle.TUPLE,
554
555 (2, 0): pickle.EMPTY_TUPLE,
556 (2, 1): pickle.TUPLE1,
557 (2, 2): pickle.TUPLE2,
558 (2, 3): pickle.TUPLE3,
559 (2, 4): pickle.TUPLE,
560 }
Guido van Rossum44f0ea52003-01-28 04:14:51 +0000561 a = ()
Guido van Rossum025bc2f2003-01-28 04:20:02 +0000562 b = (1,)
563 c = (1, 2)
564 d = (1, 2, 3)
565 e = (1, 2, 3, 4)
Tim Peters4190fb82003-02-02 16:09:05 +0000566 for proto in protocols:
Guido van Rossum44f0ea52003-01-28 04:14:51 +0000567 for x in a, b, c, d, e:
568 s = self.dumps(x, proto)
569 y = self.loads(s)
570 self.assertEqual(x, y, (proto, x, s, y))
Tim Peters1d63c9f2003-02-02 20:29:39 +0000571 expected = expected_opcode[proto, len(x)]
Tim Peters22e71712003-02-03 22:27:38 +0000572 self.assertEqual(opcode_in_pickle(expected, s), True)
Tim Peters1d63c9f2003-02-02 20:29:39 +0000573
Guido van Rossum7d97d312003-01-28 04:25:27 +0000574 def test_singletons(self):
Tim Peters61bf2572003-02-03 21:31:22 +0000575 # Map (proto, singleton) to expected opcode.
576 expected_opcode = {(0, None): pickle.NONE,
577 (1, None): pickle.NONE,
578 (2, None): pickle.NONE,
579
580 (0, True): pickle.INT,
581 (1, True): pickle.INT,
582 (2, True): pickle.NEWTRUE,
583
584 (0, False): pickle.INT,
585 (1, False): pickle.INT,
586 (2, False): pickle.NEWFALSE,
587 }
Tim Peters4190fb82003-02-02 16:09:05 +0000588 for proto in protocols:
Guido van Rossum7d97d312003-01-28 04:25:27 +0000589 for x in None, False, True:
590 s = self.dumps(x, proto)
591 y = self.loads(s)
592 self.assert_(x is y, (proto, x, s, y))
Tim Peters61bf2572003-02-03 21:31:22 +0000593 expected = expected_opcode[proto, x]
Tim Peters22e71712003-02-03 22:27:38 +0000594 self.assertEqual(opcode_in_pickle(expected, s), True)
Tim Peters3c67d792003-02-02 17:59:11 +0000595
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000596 def test_newobj_tuple(self):
Guido van Rossum3d8c01b2003-01-28 19:48:18 +0000597 x = MyTuple([1, 2, 3])
598 x.foo = 42
599 x.bar = "hello"
Tim Peters894453a2003-02-03 22:32:18 +0000600 for proto in protocols:
601 s = self.dumps(x, proto)
602 y = self.loads(s)
603 self.assertEqual(tuple(x), tuple(y))
604 self.assertEqual(x.__dict__, y.__dict__)
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000605
606 def test_newobj_list(self):
Guido van Rossum3d8c01b2003-01-28 19:48:18 +0000607 x = MyList([1, 2, 3])
608 x.foo = 42
609 x.bar = "hello"
Tim Peters894453a2003-02-03 22:32:18 +0000610 for proto in protocols:
611 s = self.dumps(x, proto)
612 y = self.loads(s)
613 self.assertEqual(list(x), list(y))
614 self.assertEqual(x.__dict__, y.__dict__)
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000615
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000616 def test_newobj_generic(self):
Tim Peters5013bd92003-02-03 22:28:41 +0000617 for proto in protocols:
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000618 for C in myclasses:
619 B = C.__base__
620 x = C(C.sample)
621 x.foo = 42
622 s = self.dumps(x, proto)
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000623 y = self.loads(s)
624 detail = (proto, C, B, x, y, type(y))
625 self.assertEqual(B(x), B(y), detail)
626 self.assertEqual(x.__dict__, y.__dict__, detail)
627
Tim Peters22e71712003-02-03 22:27:38 +0000628 # Register a type with copy_reg, with extension code extcode. Pickle
629 # an object of that type. Check that the resulting pickle uses opcode
630 # (EXT[124]) under proto 2, and not in proto 1.
Tim Peters3e667d52003-02-04 21:47:44 +0000631
Tim Peters22e71712003-02-03 22:27:38 +0000632 def produce_global_ext(self, extcode, opcode):
Tim Peters3e667d52003-02-04 21:47:44 +0000633 e = ExtensionSaver(extcode)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000634 try:
Tim Peters3e667d52003-02-04 21:47:44 +0000635 copy_reg.add_extension(__name__, "MyList", extcode)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000636 x = MyList([1, 2, 3])
637 x.foo = 42
638 x.bar = "hello"
639
Tim Peters22e71712003-02-03 22:27:38 +0000640 # Dump using protocol 1 for comparison.
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000641 s1 = self.dumps(x, 1)
Tim Peters3e667d52003-02-04 21:47:44 +0000642 self.assert_(__name__ in s1)
643 self.assert_("MyList" in s1)
644 self.assertEqual(opcode_in_pickle(opcode, s1), False)
645
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000646 y = self.loads(s1)
647 self.assertEqual(list(x), list(y))
648 self.assertEqual(x.__dict__, y.__dict__)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000649
Tim Peters22e71712003-02-03 22:27:38 +0000650 # Dump using protocol 2 for test.
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000651 s2 = self.dumps(x, 2)
Tim Peters3e667d52003-02-04 21:47:44 +0000652 self.assert_(__name__ not in s2)
653 self.assert_("MyList" not in s2)
654 self.assertEqual(opcode_in_pickle(opcode, s2), True)
655
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000656 y = self.loads(s2)
657 self.assertEqual(list(x), list(y))
658 self.assertEqual(x.__dict__, y.__dict__)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000659
660 finally:
Tim Peters3e667d52003-02-04 21:47:44 +0000661 e.restore()
Tim Peters22e71712003-02-03 22:27:38 +0000662
663 def test_global_ext1(self):
Tim Peters3e667d52003-02-04 21:47:44 +0000664 self.produce_global_ext(0x00000001, pickle.EXT1) # smallest EXT1 code
665 self.produce_global_ext(0x000000ff, pickle.EXT1) # largest EXT1 code
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000666
667 def test_global_ext2(self):
Tim Peters3e667d52003-02-04 21:47:44 +0000668 self.produce_global_ext(0x00000100, pickle.EXT2) # smallest EXT2 code
669 self.produce_global_ext(0x0000ffff, pickle.EXT2) # largest EXT2 code
670 self.produce_global_ext(0x0000abcd, pickle.EXT2) # check endianness
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000671
672 def test_global_ext4(self):
Tim Peters3e667d52003-02-04 21:47:44 +0000673 self.produce_global_ext(0x00010000, pickle.EXT4) # smallest EXT4 code
674 self.produce_global_ext(0x7fffffff, pickle.EXT4) # largest EXT4 code
675 self.produce_global_ext(0x12abcdef, pickle.EXT4) # check endianness
676
Tim Peters8d2613a2003-02-11 16:40:16 +0000677 def test_list_chunking(self):
678 n = 10 # too small to chunk
679 x = range(n)
680 for proto in protocols:
681 s = self.dumps(x, proto)
682 y = self.loads(s)
683 self.assertEqual(x, y)
684 num_appends = count_opcode(pickle.APPENDS, s)
685 self.assertEqual(num_appends, proto > 0)
686
687 n = 2500 # expect at least two chunks when proto > 0
688 x = range(n)
689 for proto in protocols:
690 s = self.dumps(x, proto)
691 y = self.loads(s)
692 self.assertEqual(x, y)
693 num_appends = count_opcode(pickle.APPENDS, s)
694 if proto == 0:
695 self.assertEqual(num_appends, 0)
696 else:
697 self.failUnless(num_appends >= 2)
698
699 def test_dict_chunking(self):
700 n = 10 # too small to chunk
701 x = dict.fromkeys(range(n))
702 for proto in protocols:
703 s = self.dumps(x, proto)
704 y = self.loads(s)
705 self.assertEqual(x, y)
706 num_setitems = count_opcode(pickle.SETITEMS, s)
707 self.assertEqual(num_setitems, proto > 0)
708
709 n = 2500 # expect at least two chunks when proto > 0
710 x = dict.fromkeys(range(n))
711 for proto in protocols:
712 s = self.dumps(x, proto)
713 y = self.loads(s)
714 self.assertEqual(x, y)
715 num_setitems = count_opcode(pickle.SETITEMS, s)
716 if proto == 0:
717 self.assertEqual(num_setitems, 0)
718 else:
719 self.failUnless(num_setitems >= 2)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000720
Tim Peterse9ef2032003-02-13 18:42:00 +0000721 def test_simple_newobj(self):
722 x = object.__new__(SimpleNewObj) # avoid __init__
723 x.abc = 666
724 for proto in protocols:
725 s = self.dumps(x, proto)
726 self.assertEqual(opcode_in_pickle(pickle.NEWOBJ, s), proto >= 2)
727 y = self.loads(s) # will raise TypeError if __init__ called
728 self.assertEqual(y.abc, 666)
729 self.assertEqual(x.__dict__, y.__dict__)
730
Tim Peters42f08ac2003-02-11 22:43:24 +0000731 def test_newobj_list_slots(self):
732 x = SlotList([1, 2, 3])
733 x.foo = 42
734 x.bar = "hello"
735 s = self.dumps(x, 2)
736 y = self.loads(s)
737 self.assertEqual(list(x), list(y))
738 self.assertEqual(x.__dict__, y.__dict__)
739 self.assertEqual(x.foo, y.foo)
740 self.assertEqual(x.bar, y.bar)
741
Tim Peters080c88b2003-02-15 03:01:11 +0000742
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000743class MyInt(int):
744 sample = 1
745
746class MyLong(long):
747 sample = 1L
748
749class MyFloat(float):
750 sample = 1.0
751
752class MyComplex(complex):
753 sample = 1.0 + 0.0j
754
755class MyStr(str):
756 sample = "hello"
757
758class MyUnicode(unicode):
759 sample = u"hello \u1234"
760
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000761class MyTuple(tuple):
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000762 sample = (1, 2, 3)
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000763
764class MyList(list):
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000765 sample = [1, 2, 3]
766
767class MyDict(dict):
768 sample = {"a": 1, "b": 2}
769
770myclasses = [MyInt, MyLong, MyFloat,
771 # MyComplex, # XXX complex somehow doesn't work here :-(
772 MyStr, MyUnicode,
773 MyTuple, MyList, MyDict]
774
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000775
Guido van Rossumc8d6ef52003-01-28 22:02:31 +0000776class SlotList(MyList):
777 __slots__ = ["foo"]
778
Tim Peterse9ef2032003-02-13 18:42:00 +0000779class SimpleNewObj(object):
780 def __init__(self, a, b, c):
781 # raise an error, to make sure this isn't called
782 raise TypeError("SimpleNewObj.__init__() didn't expect to get called")
783
Jeremy Hylton66426532001-10-15 21:38:56 +0000784class AbstractPickleModuleTests(unittest.TestCase):
785
786 def test_dump_closed_file(self):
Guido van Rossum3b0a3292002-08-09 16:38:32 +0000787 import os
788 f = open(TESTFN, "w")
789 try:
790 f.close()
791 self.assertRaises(ValueError, self.module.dump, 123, f)
792 finally:
793 os.remove(TESTFN)
Jeremy Hylton66426532001-10-15 21:38:56 +0000794
795 def test_load_closed_file(self):
Guido van Rossum3b0a3292002-08-09 16:38:32 +0000796 import os
797 f = open(TESTFN, "w")
798 try:
799 f.close()
800 self.assertRaises(ValueError, self.module.dump, 123, f)
801 finally:
802 os.remove(TESTFN)
Jeremy Hylton4c8be852002-11-13 22:10:47 +0000803
Tim Petersc0c93702003-02-13 19:30:57 +0000804 def test_highest_protocol(self):
805 # Of course this needs to be changed when HIGHEST_PROTOCOL changes.
806 self.assertEqual(self.module.HIGHEST_PROTOCOL, 2)
807
808
Jeremy Hylton4c8be852002-11-13 22:10:47 +0000809class AbstractPersistentPicklerTests(unittest.TestCase):
810
811 # This class defines persistent_id() and persistent_load()
812 # functions that should be used by the pickler. All even integers
813 # are pickled using persistent ids.
814
815 def persistent_id(self, object):
816 if isinstance(object, int) and object % 2 == 0:
817 self.id_count += 1
818 return str(object)
819 else:
820 return None
821
822 def persistent_load(self, oid):
823 self.load_count += 1
824 object = int(oid)
825 assert object % 2 == 0
826 return object
827
828 def test_persistence(self):
829 self.id_count = 0
830 self.load_count = 0
831 L = range(10)
832 self.assertEqual(self.loads(self.dumps(L)), L)
833 self.assertEqual(self.id_count, 5)
834 self.assertEqual(self.load_count, 5)
835
836 def test_bin_persistence(self):
837 self.id_count = 0
838 self.load_count = 0
839 L = range(10)
840 self.assertEqual(self.loads(self.dumps(L, 1)), L)
841 self.assertEqual(self.id_count, 5)
842 self.assertEqual(self.load_count, 5)