blob: 734f2a30c0b2ded73b3eb8dc06481b5a8e7320bc [file] [log] [blame]
Jeremy Hylton66426532001-10-15 21:38:56 +00001import unittest
Tim Peters4190fb82003-02-02 16:09:05 +00002import pickle
Tim Peters31f119e2003-02-03 16:20:13 +00003import pickletools
Tim Peters3e667d52003-02-04 21:47:44 +00004import copy_reg
Tim Peters4190fb82003-02-02 16:09:05 +00005
Guido van Rossum3b0a3292002-08-09 16:38:32 +00006from test.test_support import TestFailed, have_unicode, TESTFN
Tim Peterse089c682001-04-10 03:41:41 +00007
Tim Petersee1a53c2003-02-02 02:57:53 +00008# Tests that try a number of pickle protocols should have a
9# for proto in protocols:
10# kind of outer loop. Bump the 3 to 4 if/when protocol 3 is invented.
11protocols = range(3)
12
Tim Peters22e71712003-02-03 22:27:38 +000013
14# Return True if opcode code appears in the pickle, else False.
15def opcode_in_pickle(code, pickle):
16 for op, dummy, dummy in pickletools.genops(pickle):
17 if op.code == code:
18 return True
19 return False
20
Tim Peters8d2613a2003-02-11 16:40:16 +000021# Return the number of times opcode code appears in pickle.
22def count_opcode(code, pickle):
23 n = 0
24 for op, dummy, dummy in pickletools.genops(pickle):
25 if op.code == code:
26 n += 1
27 return n
28
Tim Peters3e667d52003-02-04 21:47:44 +000029# We can't very well test the extension registry without putting known stuff
30# in it, but we have to be careful to restore its original state. Code
31# should do this:
32#
33# e = ExtensionSaver(extension_code)
34# try:
35# fiddle w/ the extension registry's stuff for extension_code
36# finally:
37# e.restore()
38
39class ExtensionSaver:
40 # Remember current registration for code (if any), and remove it (if
41 # there is one).
42 def __init__(self, code):
43 self.code = code
44 if code in copy_reg._inverted_registry:
45 self.pair = copy_reg._inverted_registry[code]
46 copy_reg.remove_extension(self.pair[0], self.pair[1], code)
47 else:
48 self.pair = None
49
50 # Restore previous registration for code.
51 def restore(self):
52 code = self.code
53 curpair = copy_reg._inverted_registry.get(code)
54 if curpair is not None:
55 copy_reg.remove_extension(curpair[0], curpair[1], code)
56 pair = self.pair
57 if pair is not None:
58 copy_reg.add_extension(pair[0], pair[1], code)
59
Jeremy Hylton66426532001-10-15 21:38:56 +000060class C:
61 def __cmp__(self, other):
62 return cmp(self.__dict__, other.__dict__)
63
64import __main__
65__main__.C = C
66C.__module__ = "__main__"
67
68class myint(int):
69 def __init__(self, x):
70 self.str = str(x)
71
72class initarg(C):
Guido van Rossum1444f672001-12-19 16:38:29 +000073
Jeremy Hylton66426532001-10-15 21:38:56 +000074 def __init__(self, a, b):
75 self.a = a
76 self.b = b
77
78 def __getinitargs__(self):
79 return self.a, self.b
80
Guido van Rossum04a86612001-12-19 16:58:54 +000081class metaclass(type):
82 pass
83
84class use_metaclass(object):
85 __metaclass__ = metaclass
86
Tim Peters70b02d72003-02-02 17:26:40 +000087# DATA0 .. DATA2 are the pickles we expect under the various protocols, for
88# the object returned by create_data().
89# XXX DATA2 doesn't exist yet, as it's not fully implemented in cPickle.
Tim Petersee1a53c2003-02-02 02:57:53 +000090
Jeremy Hylton66426532001-10-15 21:38:56 +000091# break into multiple strings to avoid confusing font-lock-mode
Tim Peters70b02d72003-02-02 17:26:40 +000092DATA0 = """(lp1
Tim Peterse9358162001-01-22 22:05:20 +000093I0
94aL1L
Tim Peters461922a2001-04-09 20:07:05 +000095aF2
Tim Peterse9358162001-01-22 22:05:20 +000096ac__builtin__
97complex
Tim Peters461922a2001-04-09 20:07:05 +000098p2
99""" + \
100"""(F3
101F0
102tRp3
103aI1
104aI-1
105aI255
106aI-255
107aI-256
108aI65535
109aI-65535
110aI-65536
111aI2147483647
112aI-2147483647
113aI-2147483648
114a""" + \
115"""(S'abc'
Tim Peterse9358162001-01-22 22:05:20 +0000116p4
117g4
Tim Peters461922a2001-04-09 20:07:05 +0000118""" + \
Guido van Rossum42f92da2001-04-16 00:28:21 +0000119"""(i__main__
Tim Peterse9358162001-01-22 22:05:20 +0000120C
121p5
Tim Peters461922a2001-04-09 20:07:05 +0000122""" + \
Tim Peterse9358162001-01-22 22:05:20 +0000123"""(dp6
124S'foo'
125p7
126I1
127sS'bar'
128p8
129I2
130sbg5
131tp9
132ag9
133aI5
134a.
135"""
136
Tim Peters70b02d72003-02-02 17:26:40 +0000137# Disassembly of DATA0.
138DATA0_DIS = """\
139 0: ( MARK
140 1: l LIST (MARK at 0)
141 2: p PUT 1
142 5: I INT 0
143 8: a APPEND
144 9: L LONG 1L
145 13: a APPEND
146 14: F FLOAT 2.0
147 17: a APPEND
148 18: c GLOBAL '__builtin__ complex'
149 39: p PUT 2
150 42: ( MARK
151 43: F FLOAT 3.0
152 46: F FLOAT 0.0
153 49: t TUPLE (MARK at 42)
154 50: R REDUCE
155 51: p PUT 3
156 54: a APPEND
157 55: I INT 1
158 58: a APPEND
159 59: I INT -1
160 63: a APPEND
161 64: I INT 255
162 69: a APPEND
163 70: I INT -255
164 76: a APPEND
165 77: I INT -256
166 83: a APPEND
167 84: I INT 65535
168 91: a APPEND
169 92: I INT -65535
170 100: a APPEND
171 101: I INT -65536
172 109: a APPEND
173 110: I INT 2147483647
174 122: a APPEND
175 123: I INT -2147483647
176 136: a APPEND
177 137: I INT -2147483648
178 150: a APPEND
179 151: ( MARK
180 152: S STRING 'abc'
181 159: p PUT 4
182 162: g GET 4
183 165: ( MARK
184 166: i INST '__main__ C' (MARK at 165)
185 178: p PUT 5
186 181: ( MARK
187 182: d DICT (MARK at 181)
188 183: p PUT 6
189 186: S STRING 'foo'
190 193: p PUT 7
191 196: I INT 1
192 199: s SETITEM
193 200: S STRING 'bar'
194 207: p PUT 8
195 210: I INT 2
196 213: s SETITEM
197 214: b BUILD
198 215: g GET 5
199 218: t TUPLE (MARK at 151)
200 219: p PUT 9
201 222: a APPEND
202 223: g GET 9
203 226: a APPEND
204 227: I INT 5
205 230: a APPEND
206 231: . STOP
207highest protocol among opcodes = 0
208"""
209
210DATA1 = (']q\x01(K\x00L1L\nG@\x00\x00\x00\x00\x00\x00\x00'
211 'c__builtin__\ncomplex\nq\x02(G@\x08\x00\x00\x00\x00\x00'
212 '\x00G\x00\x00\x00\x00\x00\x00\x00\x00tRq\x03K\x01J\xff\xff'
213 '\xff\xffK\xffJ\x01\xff\xff\xffJ\x00\xff\xff\xffM\xff\xff'
214 'J\x01\x00\xff\xffJ\x00\x00\xff\xffJ\xff\xff\xff\x7fJ\x01\x00'
215 '\x00\x80J\x00\x00\x00\x80(U\x03abcq\x04h\x04(c__main__\n'
216 'C\nq\x05oq\x06}q\x07(U\x03fooq\x08K\x01U\x03barq\tK\x02ubh'
217 '\x06tq\nh\nK\x05e.'
218 )
219
220# Disassembly of DATA1.
221DATA1_DIS = """\
222 0: ] EMPTY_LIST
223 1: q BINPUT 1
224 3: ( MARK
225 4: K BININT1 0
226 6: L LONG 1L
227 10: G BINFLOAT 2.0
228 19: c GLOBAL '__builtin__ complex'
229 40: q BINPUT 2
230 42: ( MARK
231 43: G BINFLOAT 3.0
232 52: G BINFLOAT 0.0
233 61: t TUPLE (MARK at 42)
234 62: R REDUCE
235 63: q BINPUT 3
236 65: K BININT1 1
237 67: J BININT -1
238 72: K BININT1 255
239 74: J BININT -255
240 79: J BININT -256
241 84: M BININT2 65535
242 87: J BININT -65535
243 92: J BININT -65536
244 97: J BININT 2147483647
245 102: J BININT -2147483647
246 107: J BININT -2147483648
247 112: ( MARK
248 113: U SHORT_BINSTRING 'abc'
249 118: q BINPUT 4
250 120: h BINGET 4
251 122: ( MARK
252 123: c GLOBAL '__main__ C'
253 135: q BINPUT 5
254 137: o OBJ (MARK at 122)
255 138: q BINPUT 6
256 140: } EMPTY_DICT
257 141: q BINPUT 7
258 143: ( MARK
259 144: U SHORT_BINSTRING 'foo'
260 149: q BINPUT 8
261 151: K BININT1 1
262 153: U SHORT_BINSTRING 'bar'
263 158: q BINPUT 9
264 160: K BININT1 2
265 162: u SETITEMS (MARK at 143)
266 163: b BUILD
267 164: h BINGET 6
268 166: t TUPLE (MARK at 112)
269 167: q BINPUT 10
270 169: h BINGET 10
271 171: K BININT1 5
272 173: e APPENDS (MARK at 3)
273 174: . STOP
274highest protocol among opcodes = 1
275"""
Tim Peterse0c446b2001-10-18 21:57:37 +0000276
Jeremy Hylton66426532001-10-15 21:38:56 +0000277def create_data():
Tim Peterse9358162001-01-22 22:05:20 +0000278 c = C()
279 c.foo = 1
280 c.bar = 2
281 x = [0, 1L, 2.0, 3.0+0j]
Tim Peters461922a2001-04-09 20:07:05 +0000282 # Append some integer test cases at cPickle.c's internal size
283 # cutoffs.
284 uint1max = 0xff
285 uint2max = 0xffff
286 int4max = 0x7fffffff
287 x.extend([1, -1,
288 uint1max, -uint1max, -uint1max-1,
289 uint2max, -uint2max, -uint2max-1,
290 int4max, -int4max, -int4max-1])
Tim Peterse9358162001-01-22 22:05:20 +0000291 y = ('abc', 'abc', c, c)
292 x.append(y)
293 x.append(y)
294 x.append(5)
Jeremy Hylton66426532001-10-15 21:38:56 +0000295 return x
Tim Petersc58440f2001-04-09 17:16:31 +0000296
Jeremy Hylton66426532001-10-15 21:38:56 +0000297class AbstractPickleTests(unittest.TestCase):
Tim Peters70b02d72003-02-02 17:26:40 +0000298 # Subclass must define self.dumps, self.loads, self.error.
Tim Petersc58440f2001-04-09 17:16:31 +0000299
Jeremy Hylton66426532001-10-15 21:38:56 +0000300 _testdata = create_data()
Tim Petersc58440f2001-04-09 17:16:31 +0000301
Jeremy Hylton66426532001-10-15 21:38:56 +0000302 def setUp(self):
Tim Peterse9358162001-01-22 22:05:20 +0000303 pass
Tim Petersc58440f2001-04-09 17:16:31 +0000304
Jeremy Hylton66426532001-10-15 21:38:56 +0000305 def test_misc(self):
306 # test various datatypes not tested by testdata
Tim Peters70b02d72003-02-02 17:26:40 +0000307 for proto in protocols:
308 x = myint(4)
309 s = self.dumps(x, proto)
310 y = self.loads(s)
311 self.assertEqual(x, y)
Tim Peterse9358162001-01-22 22:05:20 +0000312
Tim Peters70b02d72003-02-02 17:26:40 +0000313 x = (1, ())
314 s = self.dumps(x, proto)
315 y = self.loads(s)
316 self.assertEqual(x, y)
Tim Peterse9358162001-01-22 22:05:20 +0000317
Tim Peters70b02d72003-02-02 17:26:40 +0000318 x = initarg(1, x)
319 s = self.dumps(x, proto)
320 y = self.loads(s)
321 self.assertEqual(x, y)
Tim Peterse9358162001-01-22 22:05:20 +0000322
Jeremy Hylton66426532001-10-15 21:38:56 +0000323 # XXX test __reduce__ protocol?
324
Tim Peters70b02d72003-02-02 17:26:40 +0000325 def test_roundtrip_equality(self):
326 expected = self._testdata
327 for proto in protocols:
328 s = self.dumps(expected, proto)
329 got = self.loads(s)
330 self.assertEqual(expected, got)
Jeremy Hylton66426532001-10-15 21:38:56 +0000331
Tim Peters70b02d72003-02-02 17:26:40 +0000332 def test_load_from_canned_string(self):
333 expected = self._testdata
334 for canned in DATA0, DATA1:
335 got = self.loads(canned)
336 self.assertEqual(expected, got)
Jeremy Hylton66426532001-10-15 21:38:56 +0000337
Tim Peters70b02d72003-02-02 17:26:40 +0000338 # There are gratuitous differences between pickles produced by
339 # pickle and cPickle, largely because cPickle starts PUT indices at
340 # 1 and pickle starts them at 0. See XXX comment in cPickle's put2() --
341 # there's a comment with an exclamation point there whose meaning
342 # is a mystery. cPickle also suppresses PUT for objects with a refcount
343 # of 1.
344 def dont_test_disassembly(self):
345 from cStringIO import StringIO
346 from pickletools import dis
347
348 for proto, expected in (0, DATA0_DIS), (1, DATA1_DIS):
349 s = self.dumps(self._testdata, proto)
350 filelike = StringIO()
351 dis(s, out=filelike)
352 got = filelike.getvalue()
353 self.assertEqual(expected, got)
Jeremy Hylton66426532001-10-15 21:38:56 +0000354
355 def test_recursive_list(self):
356 l = []
357 l.append(l)
Tim Peters70b02d72003-02-02 17:26:40 +0000358 for proto in protocols:
359 s = self.dumps(l, proto)
360 x = self.loads(s)
361 self.assertEqual(x, l)
362 self.assertEqual(x, x[0])
363 self.assertEqual(id(x), id(x[0]))
Jeremy Hylton66426532001-10-15 21:38:56 +0000364
365 def test_recursive_dict(self):
366 d = {}
367 d[1] = d
Tim Peters70b02d72003-02-02 17:26:40 +0000368 for proto in protocols:
369 s = self.dumps(d, proto)
370 x = self.loads(s)
371 self.assertEqual(x, d)
372 self.assertEqual(x[1], x)
373 self.assertEqual(id(x[1]), id(x))
Jeremy Hylton66426532001-10-15 21:38:56 +0000374
375 def test_recursive_inst(self):
376 i = C()
377 i.attr = i
Tim Peters70b02d72003-02-02 17:26:40 +0000378 for proto in protocols:
379 s = self.dumps(i, 2)
380 x = self.loads(s)
381 self.assertEqual(x, i)
382 self.assertEqual(x.attr, x)
383 self.assertEqual(id(x.attr), id(x))
Jeremy Hylton66426532001-10-15 21:38:56 +0000384
385 def test_recursive_multi(self):
386 l = []
387 d = {1:l}
388 i = C()
389 i.attr = d
390 l.append(i)
Tim Peters70b02d72003-02-02 17:26:40 +0000391 for proto in protocols:
392 s = self.dumps(l, proto)
393 x = self.loads(s)
394 self.assertEqual(x, l)
395 self.assertEqual(x[0], i)
396 self.assertEqual(x[0].attr, d)
397 self.assertEqual(x[0].attr[1], x)
398 self.assertEqual(x[0].attr[1][0], i)
399 self.assertEqual(x[0].attr[1][0].attr, d)
Jeremy Hylton66426532001-10-15 21:38:56 +0000400
401 def test_garyp(self):
402 self.assertRaises(self.error, self.loads, 'garyp')
403
404 def test_insecure_strings(self):
405 insecure = ["abc", "2 + 2", # not quoted
Martin v. Löwis8a8da792002-08-14 07:46:28 +0000406 #"'abc' + 'def'", # not a single quoted string
Jeremy Hylton66426532001-10-15 21:38:56 +0000407 "'abc", # quote is not closed
408 "'abc\"", # open quote and close quote don't match
409 "'abc' ?", # junk after close quote
Martin v. Löwiseb3f00a2002-08-14 08:22:50 +0000410 "'\\'", # trailing backslash
Jeremy Hylton66426532001-10-15 21:38:56 +0000411 # some tests of the quoting rules
Martin v. Löwis8a8da792002-08-14 07:46:28 +0000412 #"'abc\"\''",
413 #"'\\\\a\'\'\'\\\'\\\\\''",
Jeremy Hylton66426532001-10-15 21:38:56 +0000414 ]
415 for s in insecure:
416 buf = "S" + s + "\012p0\012."
417 self.assertRaises(ValueError, self.loads, buf)
418
Martin v. Löwis339d0f72001-08-17 18:39:25 +0000419 if have_unicode:
Jeremy Hylton66426532001-10-15 21:38:56 +0000420 def test_unicode(self):
421 endcases = [unicode(''), unicode('<\\u>'), unicode('<\\\u1234>'),
422 unicode('<\n>'), unicode('<\\>')]
Tim Petersee1a53c2003-02-02 02:57:53 +0000423 for proto in protocols:
424 for u in endcases:
425 p = self.dumps(u, proto)
426 u2 = self.loads(p)
427 self.assertEqual(u2, u)
Tim Peterse089c682001-04-10 03:41:41 +0000428
Jeremy Hylton66426532001-10-15 21:38:56 +0000429 def test_ints(self):
430 import sys
Tim Petersee1a53c2003-02-02 02:57:53 +0000431 for proto in protocols:
432 n = sys.maxint
433 while n:
434 for expected in (-n, n):
435 s = self.dumps(expected, proto)
436 n2 = self.loads(s)
437 self.assertEqual(expected, n2)
438 n = n >> 1
Tim Peters19ef62d2001-08-28 22:21:18 +0000439
Jeremy Hylton66426532001-10-15 21:38:56 +0000440 def test_maxint64(self):
441 maxint64 = (1L << 63) - 1
442 data = 'I' + str(maxint64) + '\n.'
443 got = self.loads(data)
444 self.assertEqual(got, maxint64)
445
446 # Try too with a bogus literal.
447 data = 'I' + str(maxint64) + 'JUNK\n.'
448 self.assertRaises(ValueError, self.loads, data)
449
Tim Petersee1a53c2003-02-02 02:57:53 +0000450 def test_long(self):
451 for proto in protocols:
Tim Petersbf2674b2003-02-02 07:51:32 +0000452 # 256 bytes is where LONG4 begins.
Tim Petersee1a53c2003-02-02 02:57:53 +0000453 for nbits in 1, 8, 8*254, 8*255, 8*256, 8*257:
454 nbase = 1L << nbits
455 for npos in nbase-1, nbase, nbase+1:
456 for n in npos, -npos:
457 pickle = self.dumps(n, proto)
458 got = self.loads(pickle)
459 self.assertEqual(n, got)
460 # Try a monster. This is quadratic-time in protos 0 & 1, so don't
461 # bother with those.
Tim Petersee1a53c2003-02-02 02:57:53 +0000462 nbase = long("deadbeeffeedface", 16)
463 nbase += nbase << 1000000
464 for n in nbase, -nbase:
Tim Petersee1a53c2003-02-02 02:57:53 +0000465 p = self.dumps(n, 2)
Tim Petersee1a53c2003-02-02 02:57:53 +0000466 got = self.loads(p)
Tim Petersee1a53c2003-02-02 02:57:53 +0000467 self.assertEqual(n, got)
468
Jeremy Hylton66426532001-10-15 21:38:56 +0000469 def test_reduce(self):
Tim Peters19ef62d2001-08-28 22:21:18 +0000470 pass
Jeremy Hylton66426532001-10-15 21:38:56 +0000471
472 def test_getinitargs(self):
473 pass
474
Guido van Rossum04a86612001-12-19 16:58:54 +0000475 def test_metaclass(self):
476 a = use_metaclass()
Tim Peters70b02d72003-02-02 17:26:40 +0000477 for proto in protocols:
478 s = self.dumps(a, proto)
479 b = self.loads(s)
480 self.assertEqual(a.__class__, b.__class__)
Guido van Rossum04a86612001-12-19 16:58:54 +0000481
Michael W. Hudson7bb466a2002-03-05 13:27:58 +0000482 def test_structseq(self):
483 import time
Michael W. Hudson0e025302002-03-06 17:11:18 +0000484 import os
Tim Peters70b02d72003-02-02 17:26:40 +0000485
486 t = time.localtime()
487 for proto in protocols:
488 s = self.dumps(t, proto)
Michael W. Hudson0e025302002-03-06 17:11:18 +0000489 u = self.loads(s)
490 self.assertEqual(t, u)
Tim Peters70b02d72003-02-02 17:26:40 +0000491 if hasattr(os, "stat"):
492 t = os.stat(os.curdir)
493 s = self.dumps(t, proto)
494 u = self.loads(s)
495 self.assertEqual(t, u)
496 if hasattr(os, "statvfs"):
497 t = os.statvfs(os.curdir)
498 s = self.dumps(t, proto)
499 u = self.loads(s)
500 self.assertEqual(t, u)
Michael W. Hudson7bb466a2002-03-05 13:27:58 +0000501
Guido van Rossumd6c9e632003-01-28 03:49:52 +0000502 # Tests for protocol 2
503
Tim Peters4190fb82003-02-02 16:09:05 +0000504 def test_proto(self):
505 build_none = pickle.NONE + pickle.STOP
506 for proto in protocols:
507 expected = build_none
508 if proto >= 2:
509 expected = pickle.PROTO + chr(proto) + expected
510 p = self.dumps(None, proto)
511 self.assertEqual(p, expected)
512
513 oob = protocols[-1] + 1 # a future protocol
514 badpickle = pickle.PROTO + chr(oob) + build_none
515 try:
516 self.loads(badpickle)
517 except ValueError, detail:
518 self.failUnless(str(detail).startswith(
519 "unsupported pickle protocol"))
520 else:
521 self.fail("expected bad protocol number to raise ValueError")
522
Guido van Rossumd6c9e632003-01-28 03:49:52 +0000523 def test_long1(self):
524 x = 12345678910111213141516178920L
Tim Peters61bf2572003-02-03 21:31:22 +0000525 for proto in protocols:
526 s = self.dumps(x, proto)
527 y = self.loads(s)
528 self.assertEqual(x, y)
Tim Peters22e71712003-02-03 22:27:38 +0000529 self.assertEqual(opcode_in_pickle(pickle.LONG1, s), proto >= 2)
Guido van Rossumd6c9e632003-01-28 03:49:52 +0000530
531 def test_long4(self):
532 x = 12345678910111213141516178920L << (256*8)
Tim Peters61bf2572003-02-03 21:31:22 +0000533 for proto in protocols:
534 s = self.dumps(x, proto)
535 y = self.loads(s)
536 self.assertEqual(x, y)
Tim Peters22e71712003-02-03 22:27:38 +0000537 self.assertEqual(opcode_in_pickle(pickle.LONG4, s), proto >= 2)
Guido van Rossumd6c9e632003-01-28 03:49:52 +0000538
Guido van Rossum44f0ea52003-01-28 04:14:51 +0000539 def test_short_tuples(self):
Tim Peters1d63c9f2003-02-02 20:29:39 +0000540 # Map (proto, len(tuple)) to expected opcode.
541 expected_opcode = {(0, 0): pickle.TUPLE,
542 (0, 1): pickle.TUPLE,
543 (0, 2): pickle.TUPLE,
544 (0, 3): pickle.TUPLE,
545 (0, 4): pickle.TUPLE,
546
547 (1, 0): pickle.EMPTY_TUPLE,
548 (1, 1): pickle.TUPLE,
549 (1, 2): pickle.TUPLE,
550 (1, 3): pickle.TUPLE,
551 (1, 4): pickle.TUPLE,
552
553 (2, 0): pickle.EMPTY_TUPLE,
554 (2, 1): pickle.TUPLE1,
555 (2, 2): pickle.TUPLE2,
556 (2, 3): pickle.TUPLE3,
557 (2, 4): pickle.TUPLE,
558 }
Guido van Rossum44f0ea52003-01-28 04:14:51 +0000559 a = ()
Guido van Rossum025bc2f2003-01-28 04:20:02 +0000560 b = (1,)
561 c = (1, 2)
562 d = (1, 2, 3)
563 e = (1, 2, 3, 4)
Tim Peters4190fb82003-02-02 16:09:05 +0000564 for proto in protocols:
Guido van Rossum44f0ea52003-01-28 04:14:51 +0000565 for x in a, b, c, d, e:
566 s = self.dumps(x, proto)
567 y = self.loads(s)
568 self.assertEqual(x, y, (proto, x, s, y))
Tim Peters1d63c9f2003-02-02 20:29:39 +0000569 expected = expected_opcode[proto, len(x)]
Tim Peters22e71712003-02-03 22:27:38 +0000570 self.assertEqual(opcode_in_pickle(expected, s), True)
Tim Peters1d63c9f2003-02-02 20:29:39 +0000571
Guido van Rossum7d97d312003-01-28 04:25:27 +0000572 def test_singletons(self):
Tim Peters61bf2572003-02-03 21:31:22 +0000573 # Map (proto, singleton) to expected opcode.
574 expected_opcode = {(0, None): pickle.NONE,
575 (1, None): pickle.NONE,
576 (2, None): pickle.NONE,
577
578 (0, True): pickle.INT,
579 (1, True): pickle.INT,
580 (2, True): pickle.NEWTRUE,
581
582 (0, False): pickle.INT,
583 (1, False): pickle.INT,
584 (2, False): pickle.NEWFALSE,
585 }
Tim Peters4190fb82003-02-02 16:09:05 +0000586 for proto in protocols:
Guido van Rossum7d97d312003-01-28 04:25:27 +0000587 for x in None, False, True:
588 s = self.dumps(x, proto)
589 y = self.loads(s)
590 self.assert_(x is y, (proto, x, s, y))
Tim Peters61bf2572003-02-03 21:31:22 +0000591 expected = expected_opcode[proto, x]
Tim Peters22e71712003-02-03 22:27:38 +0000592 self.assertEqual(opcode_in_pickle(expected, s), True)
Tim Peters3c67d792003-02-02 17:59:11 +0000593
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000594 def test_newobj_tuple(self):
Guido van Rossum3d8c01b2003-01-28 19:48:18 +0000595 x = MyTuple([1, 2, 3])
596 x.foo = 42
597 x.bar = "hello"
Tim Peters894453a2003-02-03 22:32:18 +0000598 for proto in protocols:
599 s = self.dumps(x, proto)
600 y = self.loads(s)
601 self.assertEqual(tuple(x), tuple(y))
602 self.assertEqual(x.__dict__, y.__dict__)
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000603
604 def test_newobj_list(self):
Guido van Rossum3d8c01b2003-01-28 19:48:18 +0000605 x = MyList([1, 2, 3])
606 x.foo = 42
607 x.bar = "hello"
Tim Peters894453a2003-02-03 22:32:18 +0000608 for proto in protocols:
609 s = self.dumps(x, proto)
610 y = self.loads(s)
611 self.assertEqual(list(x), list(y))
612 self.assertEqual(x.__dict__, y.__dict__)
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000613
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000614 def test_newobj_generic(self):
Tim Peters5013bd92003-02-03 22:28:41 +0000615 for proto in protocols:
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000616 for C in myclasses:
617 B = C.__base__
618 x = C(C.sample)
619 x.foo = 42
620 s = self.dumps(x, proto)
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000621 y = self.loads(s)
622 detail = (proto, C, B, x, y, type(y))
623 self.assertEqual(B(x), B(y), detail)
624 self.assertEqual(x.__dict__, y.__dict__, detail)
625
Tim Peters22e71712003-02-03 22:27:38 +0000626 # Register a type with copy_reg, with extension code extcode. Pickle
627 # an object of that type. Check that the resulting pickle uses opcode
628 # (EXT[124]) under proto 2, and not in proto 1.
Tim Peters3e667d52003-02-04 21:47:44 +0000629
Tim Peters22e71712003-02-03 22:27:38 +0000630 def produce_global_ext(self, extcode, opcode):
Tim Peters3e667d52003-02-04 21:47:44 +0000631 e = ExtensionSaver(extcode)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000632 try:
Tim Peters3e667d52003-02-04 21:47:44 +0000633 copy_reg.add_extension(__name__, "MyList", extcode)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000634 x = MyList([1, 2, 3])
635 x.foo = 42
636 x.bar = "hello"
637
Tim Peters22e71712003-02-03 22:27:38 +0000638 # Dump using protocol 1 for comparison.
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000639 s1 = self.dumps(x, 1)
Tim Peters3e667d52003-02-04 21:47:44 +0000640 self.assert_(__name__ in s1)
641 self.assert_("MyList" in s1)
642 self.assertEqual(opcode_in_pickle(opcode, s1), False)
643
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000644 y = self.loads(s1)
645 self.assertEqual(list(x), list(y))
646 self.assertEqual(x.__dict__, y.__dict__)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000647
Tim Peters22e71712003-02-03 22:27:38 +0000648 # Dump using protocol 2 for test.
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000649 s2 = self.dumps(x, 2)
Tim Peters3e667d52003-02-04 21:47:44 +0000650 self.assert_(__name__ not in s2)
651 self.assert_("MyList" not in s2)
652 self.assertEqual(opcode_in_pickle(opcode, s2), True)
653
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000654 y = self.loads(s2)
655 self.assertEqual(list(x), list(y))
656 self.assertEqual(x.__dict__, y.__dict__)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000657
658 finally:
Tim Peters3e667d52003-02-04 21:47:44 +0000659 e.restore()
Tim Peters22e71712003-02-03 22:27:38 +0000660
661 def test_global_ext1(self):
Tim Peters3e667d52003-02-04 21:47:44 +0000662 self.produce_global_ext(0x00000001, pickle.EXT1) # smallest EXT1 code
663 self.produce_global_ext(0x000000ff, pickle.EXT1) # largest EXT1 code
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000664
665 def test_global_ext2(self):
Tim Peters3e667d52003-02-04 21:47:44 +0000666 self.produce_global_ext(0x00000100, pickle.EXT2) # smallest EXT2 code
667 self.produce_global_ext(0x0000ffff, pickle.EXT2) # largest EXT2 code
668 self.produce_global_ext(0x0000abcd, pickle.EXT2) # check endianness
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000669
670 def test_global_ext4(self):
Tim Peters3e667d52003-02-04 21:47:44 +0000671 self.produce_global_ext(0x00010000, pickle.EXT4) # smallest EXT4 code
672 self.produce_global_ext(0x7fffffff, pickle.EXT4) # largest EXT4 code
673 self.produce_global_ext(0x12abcdef, pickle.EXT4) # check endianness
674
Tim Peters8d2613a2003-02-11 16:40:16 +0000675 def test_list_chunking(self):
676 n = 10 # too small to chunk
677 x = range(n)
678 for proto in protocols:
679 s = self.dumps(x, proto)
680 y = self.loads(s)
681 self.assertEqual(x, y)
682 num_appends = count_opcode(pickle.APPENDS, s)
683 self.assertEqual(num_appends, proto > 0)
684
685 n = 2500 # expect at least two chunks when proto > 0
686 x = range(n)
687 for proto in protocols:
688 s = self.dumps(x, proto)
689 y = self.loads(s)
690 self.assertEqual(x, y)
691 num_appends = count_opcode(pickle.APPENDS, s)
692 if proto == 0:
693 self.assertEqual(num_appends, 0)
694 else:
695 self.failUnless(num_appends >= 2)
696
697 def test_dict_chunking(self):
698 n = 10 # too small to chunk
699 x = dict.fromkeys(range(n))
700 for proto in protocols:
701 s = self.dumps(x, proto)
702 y = self.loads(s)
703 self.assertEqual(x, y)
704 num_setitems = count_opcode(pickle.SETITEMS, s)
705 self.assertEqual(num_setitems, proto > 0)
706
707 n = 2500 # expect at least two chunks when proto > 0
708 x = dict.fromkeys(range(n))
709 for proto in protocols:
710 s = self.dumps(x, proto)
711 y = self.loads(s)
712 self.assertEqual(x, y)
713 num_setitems = count_opcode(pickle.SETITEMS, s)
714 if proto == 0:
715 self.assertEqual(num_setitems, 0)
716 else:
717 self.failUnless(num_setitems >= 2)
Guido van Rossum0322d0f2003-01-29 06:12:46 +0000718
Tim Peters42f08ac2003-02-11 22:43:24 +0000719# XXX Temporary hack, so long as the C implementation of pickle protocol
720# XXX 2 isn't ready. When it is, move the methods in TempAbstractPickleTests
721# XXX into AbstractPickleTests above, and get rid of TempAbstractPickleTests
722# XXX along with the references to it in test_pickle.py.
723class TempAbstractPickleTests(unittest.TestCase):
724
725 def test_newobj_list_slots(self):
726 x = SlotList([1, 2, 3])
727 x.foo = 42
728 x.bar = "hello"
729 s = self.dumps(x, 2)
730 y = self.loads(s)
731 self.assertEqual(list(x), list(y))
732 self.assertEqual(x.__dict__, y.__dict__)
733 self.assertEqual(x.foo, y.foo)
734 self.assertEqual(x.bar, y.bar)
735
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000736class MyInt(int):
737 sample = 1
738
739class MyLong(long):
740 sample = 1L
741
742class MyFloat(float):
743 sample = 1.0
744
745class MyComplex(complex):
746 sample = 1.0 + 0.0j
747
748class MyStr(str):
749 sample = "hello"
750
751class MyUnicode(unicode):
752 sample = u"hello \u1234"
753
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000754class MyTuple(tuple):
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000755 sample = (1, 2, 3)
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000756
757class MyList(list):
Guido van Rossum5d9113d2003-01-29 17:58:45 +0000758 sample = [1, 2, 3]
759
760class MyDict(dict):
761 sample = {"a": 1, "b": 2}
762
763myclasses = [MyInt, MyLong, MyFloat,
764 # MyComplex, # XXX complex somehow doesn't work here :-(
765 MyStr, MyUnicode,
766 MyTuple, MyList, MyDict]
767
Guido van Rossum533dbcf2003-01-28 17:55:05 +0000768
Guido van Rossumc8d6ef52003-01-28 22:02:31 +0000769class SlotList(MyList):
770 __slots__ = ["foo"]
771
Jeremy Hylton66426532001-10-15 21:38:56 +0000772class AbstractPickleModuleTests(unittest.TestCase):
773
774 def test_dump_closed_file(self):
Guido van Rossum3b0a3292002-08-09 16:38:32 +0000775 import os
776 f = open(TESTFN, "w")
777 try:
778 f.close()
779 self.assertRaises(ValueError, self.module.dump, 123, f)
780 finally:
781 os.remove(TESTFN)
Jeremy Hylton66426532001-10-15 21:38:56 +0000782
783 def test_load_closed_file(self):
Guido van Rossum3b0a3292002-08-09 16:38:32 +0000784 import os
785 f = open(TESTFN, "w")
786 try:
787 f.close()
788 self.assertRaises(ValueError, self.module.dump, 123, f)
789 finally:
790 os.remove(TESTFN)
Jeremy Hylton4c8be852002-11-13 22:10:47 +0000791
792class AbstractPersistentPicklerTests(unittest.TestCase):
793
794 # This class defines persistent_id() and persistent_load()
795 # functions that should be used by the pickler. All even integers
796 # are pickled using persistent ids.
797
798 def persistent_id(self, object):
799 if isinstance(object, int) and object % 2 == 0:
800 self.id_count += 1
801 return str(object)
802 else:
803 return None
804
805 def persistent_load(self, oid):
806 self.load_count += 1
807 object = int(oid)
808 assert object % 2 == 0
809 return object
810
811 def test_persistence(self):
812 self.id_count = 0
813 self.load_count = 0
814 L = range(10)
815 self.assertEqual(self.loads(self.dumps(L)), L)
816 self.assertEqual(self.id_count, 5)
817 self.assertEqual(self.load_count, 5)
818
819 def test_bin_persistence(self):
820 self.id_count = 0
821 self.load_count = 0
822 L = range(10)
823 self.assertEqual(self.loads(self.dumps(L, 1)), L)
824 self.assertEqual(self.id_count, 5)
825 self.assertEqual(self.load_count, 5)