/Lib/test/pickletester.py
Python | 1370 lines | 1175 code | 116 blank | 79 comment | 77 complexity | ead9090d9a1e59a06636b3430235882e MD5 | raw file
1import unittest 2import pickle 3import cPickle 4import cStringIO 5import pickletools 6import copy_reg 7 8from test.test_support import TestFailed, have_unicode, TESTFN 9 10# Tests that try a number of pickle protocols should have a 11# for proto in protocols: 12# kind of outer loop. 13assert pickle.HIGHEST_PROTOCOL == cPickle.HIGHEST_PROTOCOL == 2 14protocols = range(pickle.HIGHEST_PROTOCOL + 1) 15 16# Copy of test.test_support.run_with_locale. This is needed to support Python 17# 2.4, which didn't include it. 18def run_with_locale(catstr, *locales): 19 def decorator(func): 20 def inner(*args, **kwds): 21 try: 22 import locale 23 category = getattr(locale, catstr) 24 orig_locale = locale.setlocale(category) 25 except AttributeError: 26 # if the test author gives us an invalid category string 27 raise 28 except: 29 # cannot retrieve original locale, so do nothing 30 locale = orig_locale = None 31 else: 32 for loc in locales: 33 try: 34 locale.setlocale(category, loc) 35 break 36 except: 37 pass 38 39 # now run the function, resetting the locale on exceptions 40 try: 41 return func(*args, **kwds) 42 finally: 43 if locale and orig_locale: 44 locale.setlocale(category, orig_locale) 45 inner.func_name = func.func_name 46 inner.__doc__ = func.__doc__ 47 return inner 48 return decorator 49 50 51# Return True if opcode code appears in the pickle, else False. 52def opcode_in_pickle(code, pickle): 53 for op, dummy, dummy in pickletools.genops(pickle): 54 if op.code == code: 55 return True 56 return False 57 58# Return the number of times opcode code appears in pickle. 59def count_opcode(code, pickle): 60 n = 0 61 for op, dummy, dummy in pickletools.genops(pickle): 62 if op.code == code: 63 n += 1 64 return n 65 66# We can't very well test the extension registry without putting known stuff 67# in it, but we have to be careful to restore its original state. Code 68# should do this: 69# 70# e = ExtensionSaver(extension_code) 71# try: 72# fiddle w/ the extension registry's stuff for extension_code 73# finally: 74# e.restore() 75 76class ExtensionSaver: 77 # Remember current registration for code (if any), and remove it (if 78 # there is one). 79 def __init__(self, code): 80 self.code = code 81 if code in copy_reg._inverted_registry: 82 self.pair = copy_reg._inverted_registry[code] 83 copy_reg.remove_extension(self.pair[0], self.pair[1], code) 84 else: 85 self.pair = None 86 87 # Restore previous registration for code. 88 def restore(self): 89 code = self.code 90 curpair = copy_reg._inverted_registry.get(code) 91 if curpair is not None: 92 copy_reg.remove_extension(curpair[0], curpair[1], code) 93 pair = self.pair 94 if pair is not None: 95 copy_reg.add_extension(pair[0], pair[1], code) 96 97class C: 98 def __cmp__(self, other): 99 return cmp(self.__dict__, other.__dict__) 100 101import __main__ 102__main__.C = C 103C.__module__ = "__main__" 104 105class myint(int): 106 def __init__(self, x): 107 self.str = str(x) 108 109class initarg(C): 110 111 def __init__(self, a, b): 112 self.a = a 113 self.b = b 114 115 def __getinitargs__(self): 116 return self.a, self.b 117 118class metaclass(type): 119 pass 120 121class use_metaclass(object): 122 __metaclass__ = metaclass 123 124# DATA0 .. DATA2 are the pickles we expect under the various protocols, for 125# the object returned by create_data(). 126 127# break into multiple strings to avoid confusing font-lock-mode 128DATA0 = """(lp1 129I0 130aL1L 131aF2 132ac__builtin__ 133complex 134p2 135""" + \ 136"""(F3 137F0 138tRp3 139aI1 140aI-1 141aI255 142aI-255 143aI-256 144aI65535 145aI-65535 146aI-65536 147aI2147483647 148aI-2147483647 149aI-2147483648 150a""" + \ 151"""(S'abc' 152p4 153g4 154""" + \ 155"""(i__main__ 156C 157p5 158""" + \ 159"""(dp6 160S'foo' 161p7 162I1 163sS'bar' 164p8 165I2 166sbg5 167tp9 168ag9 169aI5 170a. 171""" 172 173# Disassembly of DATA0. 174DATA0_DIS = """\ 175 0: ( MARK 176 1: l LIST (MARK at 0) 177 2: p PUT 1 178 5: I INT 0 179 8: a APPEND 180 9: L LONG 1L 181 13: a APPEND 182 14: F FLOAT 2.0 183 17: a APPEND 184 18: c GLOBAL '__builtin__ complex' 185 39: p PUT 2 186 42: ( MARK 187 43: F FLOAT 3.0 188 46: F FLOAT 0.0 189 49: t TUPLE (MARK at 42) 190 50: R REDUCE 191 51: p PUT 3 192 54: a APPEND 193 55: I INT 1 194 58: a APPEND 195 59: I INT -1 196 63: a APPEND 197 64: I INT 255 198 69: a APPEND 199 70: I INT -255 200 76: a APPEND 201 77: I INT -256 202 83: a APPEND 203 84: I INT 65535 204 91: a APPEND 205 92: I INT -65535 206 100: a APPEND 207 101: I INT -65536 208 109: a APPEND 209 110: I INT 2147483647 210 122: a APPEND 211 123: I INT -2147483647 212 136: a APPEND 213 137: I INT -2147483648 214 150: a APPEND 215 151: ( MARK 216 152: S STRING 'abc' 217 159: p PUT 4 218 162: g GET 4 219 165: ( MARK 220 166: i INST '__main__ C' (MARK at 165) 221 178: p PUT 5 222 181: ( MARK 223 182: d DICT (MARK at 181) 224 183: p PUT 6 225 186: S STRING 'foo' 226 193: p PUT 7 227 196: I INT 1 228 199: s SETITEM 229 200: S STRING 'bar' 230 207: p PUT 8 231 210: I INT 2 232 213: s SETITEM 233 214: b BUILD 234 215: g GET 5 235 218: t TUPLE (MARK at 151) 236 219: p PUT 9 237 222: a APPEND 238 223: g GET 9 239 226: a APPEND 240 227: I INT 5 241 230: a APPEND 242 231: . STOP 243highest protocol among opcodes = 0 244""" 245 246DATA1 = (']q\x01(K\x00L1L\nG@\x00\x00\x00\x00\x00\x00\x00' 247 'c__builtin__\ncomplex\nq\x02(G@\x08\x00\x00\x00\x00\x00' 248 '\x00G\x00\x00\x00\x00\x00\x00\x00\x00tRq\x03K\x01J\xff\xff' 249 '\xff\xffK\xffJ\x01\xff\xff\xffJ\x00\xff\xff\xffM\xff\xff' 250 'J\x01\x00\xff\xffJ\x00\x00\xff\xffJ\xff\xff\xff\x7fJ\x01\x00' 251 '\x00\x80J\x00\x00\x00\x80(U\x03abcq\x04h\x04(c__main__\n' 252 'C\nq\x05oq\x06}q\x07(U\x03fooq\x08K\x01U\x03barq\tK\x02ubh' 253 '\x06tq\nh\nK\x05e.' 254 ) 255 256# Disassembly of DATA1. 257DATA1_DIS = """\ 258 0: ] EMPTY_LIST 259 1: q BINPUT 1 260 3: ( MARK 261 4: K BININT1 0 262 6: L LONG 1L 263 10: G BINFLOAT 2.0 264 19: c GLOBAL '__builtin__ complex' 265 40: q BINPUT 2 266 42: ( MARK 267 43: G BINFLOAT 3.0 268 52: G BINFLOAT 0.0 269 61: t TUPLE (MARK at 42) 270 62: R REDUCE 271 63: q BINPUT 3 272 65: K BININT1 1 273 67: J BININT -1 274 72: K BININT1 255 275 74: J BININT -255 276 79: J BININT -256 277 84: M BININT2 65535 278 87: J BININT -65535 279 92: J BININT -65536 280 97: J BININT 2147483647 281 102: J BININT -2147483647 282 107: J BININT -2147483648 283 112: ( MARK 284 113: U SHORT_BINSTRING 'abc' 285 118: q BINPUT 4 286 120: h BINGET 4 287 122: ( MARK 288 123: c GLOBAL '__main__ C' 289 135: q BINPUT 5 290 137: o OBJ (MARK at 122) 291 138: q BINPUT 6 292 140: } EMPTY_DICT 293 141: q BINPUT 7 294 143: ( MARK 295 144: U SHORT_BINSTRING 'foo' 296 149: q BINPUT 8 297 151: K BININT1 1 298 153: U SHORT_BINSTRING 'bar' 299 158: q BINPUT 9 300 160: K BININT1 2 301 162: u SETITEMS (MARK at 143) 302 163: b BUILD 303 164: h BINGET 6 304 166: t TUPLE (MARK at 112) 305 167: q BINPUT 10 306 169: h BINGET 10 307 171: K BININT1 5 308 173: e APPENDS (MARK at 3) 309 174: . STOP 310highest protocol among opcodes = 1 311""" 312 313DATA2 = ('\x80\x02]q\x01(K\x00\x8a\x01\x01G@\x00\x00\x00\x00\x00\x00\x00' 314 'c__builtin__\ncomplex\nq\x02G@\x08\x00\x00\x00\x00\x00\x00G\x00' 315 '\x00\x00\x00\x00\x00\x00\x00\x86Rq\x03K\x01J\xff\xff\xff\xffK' 316 '\xffJ\x01\xff\xff\xffJ\x00\xff\xff\xffM\xff\xffJ\x01\x00\xff\xff' 317 'J\x00\x00\xff\xffJ\xff\xff\xff\x7fJ\x01\x00\x00\x80J\x00\x00\x00' 318 '\x80(U\x03abcq\x04h\x04(c__main__\nC\nq\x05oq\x06}q\x07(U\x03foo' 319 'q\x08K\x01U\x03barq\tK\x02ubh\x06tq\nh\nK\x05e.') 320 321# Disassembly of DATA2. 322DATA2_DIS = """\ 323 0: \x80 PROTO 2 324 2: ] EMPTY_LIST 325 3: q BINPUT 1 326 5: ( MARK 327 6: K BININT1 0 328 8: \x8a LONG1 1L 329 11: G BINFLOAT 2.0 330 20: c GLOBAL '__builtin__ complex' 331 41: q BINPUT 2 332 43: G BINFLOAT 3.0 333 52: G BINFLOAT 0.0 334 61: \x86 TUPLE2 335 62: R REDUCE 336 63: q BINPUT 3 337 65: K BININT1 1 338 67: J BININT -1 339 72: K BININT1 255 340 74: J BININT -255 341 79: J BININT -256 342 84: M BININT2 65535 343 87: J BININT -65535 344 92: J BININT -65536 345 97: J BININT 2147483647 346 102: J BININT -2147483647 347 107: J BININT -2147483648 348 112: ( MARK 349 113: U SHORT_BINSTRING 'abc' 350 118: q BINPUT 4 351 120: h BINGET 4 352 122: ( MARK 353 123: c GLOBAL '__main__ C' 354 135: q BINPUT 5 355 137: o OBJ (MARK at 122) 356 138: q BINPUT 6 357 140: } EMPTY_DICT 358 141: q BINPUT 7 359 143: ( MARK 360 144: U SHORT_BINSTRING 'foo' 361 149: q BINPUT 8 362 151: K BININT1 1 363 153: U SHORT_BINSTRING 'bar' 364 158: q BINPUT 9 365 160: K BININT1 2 366 162: u SETITEMS (MARK at 143) 367 163: b BUILD 368 164: h BINGET 6 369 166: t TUPLE (MARK at 112) 370 167: q BINPUT 10 371 169: h BINGET 10 372 171: K BININT1 5 373 173: e APPENDS (MARK at 5) 374 174: . STOP 375highest protocol among opcodes = 2 376""" 377 378def create_data(): 379 c = C() 380 c.foo = 1 381 c.bar = 2 382 x = [0, 1L, 2.0, 3.0+0j] 383 # Append some integer test cases at cPickle.c's internal size 384 # cutoffs. 385 uint1max = 0xff 386 uint2max = 0xffff 387 int4max = 0x7fffffff 388 x.extend([1, -1, 389 uint1max, -uint1max, -uint1max-1, 390 uint2max, -uint2max, -uint2max-1, 391 int4max, -int4max, -int4max-1]) 392 y = ('abc', 'abc', c, c) 393 x.append(y) 394 x.append(y) 395 x.append(5) 396 return x 397 398class AbstractPickleTests(unittest.TestCase): 399 # Subclass must define self.dumps, self.loads, self.error. 400 401 _testdata = create_data() 402 403 def setUp(self): 404 pass 405 406 def dump(self, arg, buf, proto=0): 407 data = self.dumps(arg, proto) 408 buf.write(data) 409 410 def test_misc(self): 411 # test various datatypes not tested by testdata 412 for proto in protocols: 413 x = myint(4) 414 s = self.dumps(x, proto) 415 y = self.loads(s) 416 self.assertEqual(x, y) 417 418 x = (1, ()) 419 s = self.dumps(x, proto) 420 y = self.loads(s) 421 self.assertEqual(x, y) 422 423 x = initarg(1, x) 424 s = self.dumps(x, proto) 425 y = self.loads(s) 426 self.assertEqual(x, y) 427 428 # XXX test __reduce__ protocol? 429 430 def test_roundtrip_equality(self): 431 expected = self._testdata 432 for proto in protocols: 433 s = self.dumps(expected, proto) 434 got = self.loads(s) 435 self.assertEqual(expected, got) 436 437 def test_load_from_canned_string(self): 438 expected = self._testdata 439 for canned in DATA0, DATA1, DATA2: 440 got = self.loads(canned) 441 self.assertEqual(expected, got) 442 443 # There are gratuitous differences between pickles produced by 444 # pickle and cPickle, largely because cPickle starts PUT indices at 445 # 1 and pickle starts them at 0. See XXX comment in cPickle's put2() -- 446 # there's a comment with an exclamation point there whose meaning 447 # is a mystery. cPickle also suppresses PUT for objects with a refcount 448 # of 1. 449 def dont_test_disassembly(self): 450 from pickletools import dis 451 452 for proto, expected in (0, DATA0_DIS), (1, DATA1_DIS): 453 s = self.dumps(self._testdata, proto) 454 filelike = cStringIO.StringIO() 455 dis(s, out=filelike) 456 got = filelike.getvalue() 457 self.assertEqual(expected, got) 458 459 def test_recursive_list(self): 460 l = [] 461 l.append(l) 462 for proto in protocols: 463 s = self.dumps(l, proto) 464 x = self.loads(s) 465 self.assertEqual(len(x), 1) 466 self.assert_(x is x[0]) 467 468 def test_recursive_tuple(self): 469 t = ([],) 470 t[0].append(t) 471 for proto in protocols: 472 s = self.dumps(t, proto) 473 x = self.loads(s) 474 self.assertEqual(len(x), 1) 475 self.assertEqual(len(x[0]), 1) 476 self.assert_(x is x[0][0]) 477 478 def test_recursive_dict(self): 479 d = {} 480 d[1] = d 481 for proto in protocols: 482 s = self.dumps(d, proto) 483 x = self.loads(s) 484 self.assertEqual(x.keys(), [1]) 485 self.assert_(x[1] is x) 486 487 def test_recursive_inst(self): 488 i = C() 489 i.attr = i 490 for proto in protocols: 491 s = self.dumps(i, 2) 492 x = self.loads(s) 493 self.assertEqual(dir(x), dir(i)) 494 self.assert_(x.attr is x) 495 496 def test_recursive_multi(self): 497 l = [] 498 d = {1:l} 499 i = C() 500 i.attr = d 501 l.append(i) 502 for proto in protocols: 503 s = self.dumps(l, proto) 504 x = self.loads(s) 505 self.assertEqual(len(x), 1) 506 self.assertEqual(dir(x[0]), dir(i)) 507 self.assertEqual(x[0].attr.keys(), [1]) 508 self.assert_(x[0].attr[1] is x) 509 510 def test_garyp(self): 511 self.assertRaises(self.error, self.loads, 'garyp') 512 513 def test_insecure_strings(self): 514 insecure = ["abc", "2 + 2", # not quoted 515 #"'abc' + 'def'", # not a single quoted string 516 "'abc", # quote is not closed 517 "'abc\"", # open quote and close quote don't match 518 "'abc' ?", # junk after close quote 519 "'\\'", # trailing backslash 520 # some tests of the quoting rules 521 #"'abc\"\''", 522 #"'\\\\a\'\'\'\\\'\\\\\''", 523 ] 524 for s in insecure: 525 buf = "S" + s + "\012p0\012." 526 self.assertRaises(ValueError, self.loads, buf) 527 528 if have_unicode: 529 def test_unicode(self): 530 endcases = [u'', u'<\\u>', u'<\\\u1234>', u'<\n>', 531 u'<\\>', u'<\\\U00012345>'] 532 for proto in protocols: 533 for u in endcases: 534 p = self.dumps(u, proto) 535 u2 = self.loads(p) 536 self.assertEqual(u2, u) 537 538 def test_unicode_high_plane(self): 539 t = u'\U00012345' 540 for proto in protocols: 541 p = self.dumps(t, proto) 542 t2 = self.loads(p) 543 self.assertEqual(t2, t) 544 545 def test_ints(self): 546 import sys 547 for proto in protocols: 548 n = sys.maxint 549 while n: 550 for expected in (-n, n): 551 s = self.dumps(expected, proto) 552 n2 = self.loads(s) 553 self.assertEqual(expected, n2) 554 n = n >> 1 555 556 def test_maxint64(self): 557 maxint64 = (1L << 63) - 1 558 data = 'I' + str(maxint64) + '\n.' 559 got = self.loads(data) 560 self.assertEqual(got, maxint64) 561 562 # Try too with a bogus literal. 563 data = 'I' + str(maxint64) + 'JUNK\n.' 564 self.assertRaises(ValueError, self.loads, data) 565 566 def test_long(self): 567 for proto in protocols: 568 # 256 bytes is where LONG4 begins. 569 for nbits in 1, 8, 8*254, 8*255, 8*256, 8*257: 570 nbase = 1L << nbits 571 for npos in nbase-1, nbase, nbase+1: 572 for n in npos, -npos: 573 pickle = self.dumps(n, proto) 574 got = self.loads(pickle) 575 self.assertEqual(n, got) 576 # Try a monster. This is quadratic-time in protos 0 & 1, so don't 577 # bother with those. 578 nbase = long("deadbeeffeedface", 16) 579 nbase += nbase << 1000000 580 for n in nbase, -nbase: 581 p = self.dumps(n, 2) 582 got = self.loads(p) 583 self.assertEqual(n, got) 584 585 def test_float(self): 586 test_values = [0.0, 4.94e-324, 1e-310, 7e-308, 6.626e-34, 0.1, 0.5, 587 3.14, 263.44582062374053, 6.022e23, 1e30] 588 test_values = test_values + [-x for x in test_values] 589 for proto in protocols: 590 for value in test_values: 591 pickle = self.dumps(value, proto) 592 got = self.loads(pickle) 593 self.assertEqual(value, got) 594 595 @run_with_locale('LC_ALL', 'de_DE', 'fr_FR') 596 def test_float_format(self): 597 # make sure that floats are formatted locale independent 598 self.assertEqual(self.dumps(1.2)[0:3], 'F1.') 599 600 def test_reduce(self): 601 pass 602 603 def test_getinitargs(self): 604 pass 605 606 def test_metaclass(self): 607 a = use_metaclass() 608 for proto in protocols: 609 s = self.dumps(a, proto) 610 b = self.loads(s) 611 self.assertEqual(a.__class__, b.__class__) 612 613 def test_structseq(self): 614 import time 615 import os 616 617 t = time.localtime() 618 for proto in protocols: 619 s = self.dumps(t, proto) 620 u = self.loads(s) 621 self.assertEqual(t, u) 622 if hasattr(os, "stat"): 623 t = os.stat(os.curdir) 624 s = self.dumps(t, proto) 625 u = self.loads(s) 626 self.assertEqual(t, u) 627 if hasattr(os, "statvfs"): 628 t = os.statvfs(os.curdir) 629 s = self.dumps(t, proto) 630 u = self.loads(s) 631 self.assertEqual(t, u) 632 633 # Tests for protocol 2 634 635 def test_proto(self): 636 build_none = pickle.NONE + pickle.STOP 637 for proto in protocols: 638 expected = build_none 639 if proto >= 2: 640 expected = pickle.PROTO + chr(proto) + expected 641 p = self.dumps(None, proto) 642 self.assertEqual(p, expected) 643 644 oob = protocols[-1] + 1 # a future protocol 645 badpickle = pickle.PROTO + chr(oob) + build_none 646 try: 647 self.loads(badpickle) 648 except ValueError, detail: 649 self.failUnless(str(detail).startswith( 650 "unsupported pickle protocol")) 651 else: 652 self.fail("expected bad protocol number to raise ValueError") 653 654 def test_long1(self): 655 x = 12345678910111213141516178920L 656 for proto in protocols: 657 s = self.dumps(x, proto) 658 y = self.loads(s) 659 self.assertEqual(x, y) 660 self.assertEqual(opcode_in_pickle(pickle.LONG1, s), proto >= 2) 661 662 def test_long4(self): 663 x = 12345678910111213141516178920L << (256*8) 664 for proto in protocols: 665 s = self.dumps(x, proto) 666 y = self.loads(s) 667 self.assertEqual(x, y) 668 self.assertEqual(opcode_in_pickle(pickle.LONG4, s), proto >= 2) 669 670 def test_short_tuples(self): 671 # Map (proto, len(tuple)) to expected opcode. 672 expected_opcode = {(0, 0): pickle.TUPLE, 673 (0, 1): pickle.TUPLE, 674 (0, 2): pickle.TUPLE, 675 (0, 3): pickle.TUPLE, 676 (0, 4): pickle.TUPLE, 677 678 (1, 0): pickle.EMPTY_TUPLE, 679 (1, 1): pickle.TUPLE, 680 (1, 2): pickle.TUPLE, 681 (1, 3): pickle.TUPLE, 682 (1, 4): pickle.TUPLE, 683 684 (2, 0): pickle.EMPTY_TUPLE, 685 (2, 1): pickle.TUPLE1, 686 (2, 2): pickle.TUPLE2, 687 (2, 3): pickle.TUPLE3, 688 (2, 4): pickle.TUPLE, 689 } 690 a = () 691 b = (1,) 692 c = (1, 2) 693 d = (1, 2, 3) 694 e = (1, 2, 3, 4) 695 for proto in protocols: 696 for x in a, b, c, d, e: 697 s = self.dumps(x, proto) 698 y = self.loads(s) 699 self.assertEqual(x, y, (proto, x, s, y)) 700 expected = expected_opcode[proto, len(x)] 701 self.assertEqual(opcode_in_pickle(expected, s), True) 702 703 def test_singletons(self): 704 # Map (proto, singleton) to expected opcode. 705 expected_opcode = {(0, None): pickle.NONE, 706 (1, None): pickle.NONE, 707 (2, None): pickle.NONE, 708 709 (0, True): pickle.INT, 710 (1, True): pickle.INT, 711 (2, True): pickle.NEWTRUE, 712 713 (0, False): pickle.INT, 714 (1, False): pickle.INT, 715 (2, False): pickle.NEWFALSE, 716 } 717 for proto in protocols: 718 for x in None, False, True: 719 s = self.dumps(x, proto) 720 y = self.loads(s) 721 self.assert_(x is y, (proto, x, s, y)) 722 expected = expected_opcode[proto, x] 723 self.assertEqual(opcode_in_pickle(expected, s), True) 724 725 def test_newobj_tuple(self): 726 x = MyTuple([1, 2, 3]) 727 x.foo = 42 728 x.bar = "hello" 729 for proto in protocols: 730 s = self.dumps(x, proto) 731 y = self.loads(s) 732 self.assertEqual(tuple(x), tuple(y)) 733 self.assertEqual(x.__dict__, y.__dict__) 734 735 def test_newobj_list(self): 736 x = MyList([1, 2, 3]) 737 x.foo = 42 738 x.bar = "hello" 739 for proto in protocols: 740 s = self.dumps(x, proto) 741 y = self.loads(s) 742 self.assertEqual(list(x), list(y)) 743 self.assertEqual(x.__dict__, y.__dict__) 744 745 def test_newobj_generic(self): 746 for proto in protocols: 747 for C in myclasses: 748 B = C.__base__ 749 x = C(C.sample) 750 x.foo = 42 751 s = self.dumps(x, proto) 752 y = self.loads(s) 753 detail = (proto, C, B, x, y, type(y)) 754 self.assertEqual(B(x), B(y), detail) 755 self.assertEqual(x.__dict__, y.__dict__, detail) 756 757 # Register a type with copy_reg, with extension code extcode. Pickle 758 # an object of that type. Check that the resulting pickle uses opcode 759 # (EXT[124]) under proto 2, and not in proto 1. 760 761 def produce_global_ext(self, extcode, opcode): 762 e = ExtensionSaver(extcode) 763 try: 764 copy_reg.add_extension(__name__, "MyList", extcode) 765 x = MyList([1, 2, 3]) 766 x.foo = 42 767 x.bar = "hello" 768 769 # Dump using protocol 1 for comparison. 770 s1 = self.dumps(x, 1) 771 self.assert_(__name__ in s1) 772 self.assert_("MyList" in s1) 773 self.assertEqual(opcode_in_pickle(opcode, s1), False) 774 775 y = self.loads(s1) 776 self.assertEqual(list(x), list(y)) 777 self.assertEqual(x.__dict__, y.__dict__) 778 779 # Dump using protocol 2 for test. 780 s2 = self.dumps(x, 2) 781 self.assert_(__name__ not in s2) 782 self.assert_("MyList" not in s2) 783 self.assertEqual(opcode_in_pickle(opcode, s2), True) 784 785 y = self.loads(s2) 786 self.assertEqual(list(x), list(y)) 787 self.assertEqual(x.__dict__, y.__dict__) 788 789 finally: 790 e.restore() 791 792 def test_global_ext1(self): 793 self.produce_global_ext(0x00000001, pickle.EXT1) # smallest EXT1 code 794 self.produce_global_ext(0x000000ff, pickle.EXT1) # largest EXT1 code 795 796 def test_global_ext2(self): 797 self.produce_global_ext(0x00000100, pickle.EXT2) # smallest EXT2 code 798 self.produce_global_ext(0x0000ffff, pickle.EXT2) # largest EXT2 code 799 self.produce_global_ext(0x0000abcd, pickle.EXT2) # check endianness 800 801 def test_global_ext4(self): 802 self.produce_global_ext(0x00010000, pickle.EXT4) # smallest EXT4 code 803 self.produce_global_ext(0x7fffffff, pickle.EXT4) # largest EXT4 code 804 self.produce_global_ext(0x12abcdef, pickle.EXT4) # check endianness 805 806 def test_list_chunking(self): 807 n = 10 # too small to chunk 808 x = range(n) 809 for proto in protocols: 810 s = self.dumps(x, proto) 811 y = self.loads(s) 812 self.assertEqual(x, y) 813 num_appends = count_opcode(pickle.APPENDS, s) 814 self.assertEqual(num_appends, proto > 0) 815 816 n = 2500 # expect at least two chunks when proto > 0 817 x = range(n) 818 for proto in protocols: 819 s = self.dumps(x, proto) 820 y = self.loads(s) 821 self.assertEqual(x, y) 822 num_appends = count_opcode(pickle.APPENDS, s) 823 if proto == 0: 824 self.assertEqual(num_appends, 0) 825 else: 826 self.failUnless(num_appends >= 2) 827 828 def test_dict_chunking(self): 829 n = 10 # too small to chunk 830 x = dict.fromkeys(range(n)) 831 for proto in protocols: 832 s = self.dumps(x, proto) 833 y = self.loads(s) 834 self.assertEqual(x, y) 835 num_setitems = count_opcode(pickle.SETITEMS, s) 836 self.assertEqual(num_setitems, proto > 0) 837 838 n = 2500 # expect at least two chunks when proto > 0 839 x = dict.fromkeys(range(n)) 840 for proto in protocols: 841 s = self.dumps(x, proto) 842 y = self.loads(s) 843 self.assertEqual(x, y) 844 num_setitems = count_opcode(pickle.SETITEMS, s) 845 if proto == 0: 846 self.assertEqual(num_setitems, 0) 847 else: 848 self.failUnless(num_setitems >= 2) 849 850 def test_load_empty_string(self): 851 self.assertRaises(EOFError, self.loads, "") 852 853 def test_simple_newobj(self): 854 x = object.__new__(SimpleNewObj) # avoid __init__ 855 x.abc = 666 856 for proto in protocols: 857 s = self.dumps(x, proto) 858 self.assertEqual(opcode_in_pickle(pickle.NEWOBJ, s), proto >= 2) 859 y = self.loads(s) # will raise TypeError if __init__ called 860 self.assertEqual(y.abc, 666) 861 self.assertEqual(x.__dict__, y.__dict__) 862 863 def test_newobj_list_slots(self): 864 x = SlotList([1, 2, 3]) 865 x.foo = 42 866 x.bar = "hello" 867 s = self.dumps(x, 2) 868 y = self.loads(s) 869 self.assertEqual(list(x), list(y)) 870 self.assertEqual(x.__dict__, y.__dict__) 871 self.assertEqual(x.foo, y.foo) 872 self.assertEqual(x.bar, y.bar) 873 874 def test_reduce_overrides_default_reduce_ex(self): 875 for proto in 0, 1, 2: 876 x = REX_one() 877 self.assertEqual(x._reduce_called, 0) 878 s = self.dumps(x, proto) 879 self.assertEqual(x._reduce_called, 1) 880 y = self.loads(s) 881 self.assertEqual(y._reduce_called, 0) 882 883 def test_reduce_ex_called(self): 884 for proto in 0, 1, 2: 885 x = REX_two() 886 self.assertEqual(x._proto, None) 887 s = self.dumps(x, proto) 888 self.assertEqual(x._proto, proto) 889 y = self.loads(s) 890 self.assertEqual(y._proto, None) 891 892 def test_reduce_ex_overrides_reduce(self): 893 for proto in 0, 1, 2: 894 x = REX_three() 895 self.assertEqual(x._proto, None) 896 s = self.dumps(x, proto) 897 self.assertEqual(x._proto, proto) 898 y = self.loads(s) 899 self.assertEqual(y._proto, None) 900 901 def test_reduce_ex_calls_base(self): 902 for proto in 0, 1, 2: 903 x = REX_four() 904 self.assertEqual(x._proto, None) 905 s = self.dumps(x, proto) 906 self.assertEqual(x._proto, proto) 907 y = self.loads(s) 908 self.assertEqual(y._proto, proto) 909 910 def test_reduce_calls_base(self): 911 for proto in 0, 1, 2: 912 x = REX_five() 913 self.assertEqual(x._reduce_called, 0) 914 s = self.dumps(x, proto) 915 self.assertEqual(x._reduce_called, 1) 916 y = self.loads(s) 917 self.assertEqual(y._reduce_called, 1) 918 919 def test_reduce_bad_iterator(self): 920 # Issue4176: crash when 4th and 5th items of __reduce__() 921 # are not iterators 922 class C(object): 923 def __reduce__(self): 924 # 4th item is not an iterator 925 return list, (), None, [], None 926 class D(object): 927 def __reduce__(self): 928 # 5th item is not an iterator 929 return dict, (), None, None, [] 930 931 # Protocol 0 is less strict and also accept iterables. 932 for proto in 0, 1, 2: 933 try: 934 self.dumps(C(), proto) 935 except (AttributeError, pickle.PickleError, cPickle.PickleError): 936 pass 937 try: 938 self.dumps(D(), proto) 939 except (AttributeError, pickle.PickleError, cPickle.PickleError): 940 pass 941 942 def test_many_puts_and_gets(self): 943 # Test that internal data structures correctly deal with lots of 944 # puts/gets. There were once bugs in cPickle related to this. 945 keys = ("aaa" + str(i) for i in xrange(100)) 946 large_dict = dict((k, [4, 5, 6]) for k in keys) 947 obj = [dict(large_dict), dict(large_dict), dict(large_dict)] 948 949 for proto in [0, 1, 2]: 950 dumped = self.dumps(obj, proto) 951 loaded = self.loads(dumped) 952 self.assertEqual(loaded, obj, 953 "Failed protocol %d: %r != %r" 954 % (proto, obj, loaded)) 955 956 def test_unpickle_multiple_objs_from_same_file(self): 957 data1 = ["abcdefg", "abcdefg", 44] 958 data2 = range(10) 959 f = cStringIO.StringIO() 960 self.dump(data1, f) 961 self.dump(data2, f) 962 f.seek(0) 963 964 self.assertEqual(self.load(f), data1) 965 self.assertEqual(self.load(f), data2) 966 967 968# Test classes for reduce_ex 969 970class REX_one(object): 971 _reduce_called = 0 972 def __reduce__(self): 973 self._reduce_called = 1 974 return REX_one, () 975 # No __reduce_ex__ here, but inheriting it from object 976 977class REX_two(object): 978 _proto = None 979 def __reduce_ex__(self, proto): 980 self._proto = proto 981 return REX_two, () 982 # No __reduce__ here, but inheriting it from object 983 984class REX_three(object): 985 _proto = None 986 def __reduce_ex__(self, proto): 987 self._proto = proto 988 return REX_two, () 989 def __reduce__(self): 990 raise TestFailed, "This __reduce__ shouldn't be called" 991 992class REX_four(object): 993 _proto = None 994 def __reduce_ex__(self, proto): 995 self._proto = proto 996 return object.__reduce_ex__(self, proto) 997 # Calling base class method should succeed 998 999class REX_five(object): 1000 _reduce_called = 0 1001 def __reduce__(self): 1002 self._reduce_called = 1 1003 return object.__reduce__(self) 1004 # This one used to fail with infinite recursion 1005 1006# Test classes for newobj 1007 1008class MyInt(int): 1009 sample = 1 1010 1011class MyLong(long): 1012 sample = 1L 1013 1014class MyFloat(float): 1015 sample = 1.0 1016 1017class MyComplex(complex): 1018 sample = 1.0 + 0.0j 1019 1020class MyStr(str): 1021 sample = "hello" 1022 1023class MyUnicode(unicode): 1024 sample = u"hello \u1234" 1025 1026class MyTuple(tuple): 1027 sample = (1, 2, 3) 1028 1029class MyList(list): 1030 sample = [1, 2, 3] 1031 1032class MyDict(dict): 1033 sample = {"a": 1, "b": 2} 1034 1035myclasses = [MyInt, MyLong, MyFloat, 1036 MyComplex, 1037 MyStr, MyUnicode, 1038 MyTuple, MyList, MyDict] 1039 1040 1041class SlotList(MyList): 1042 __slots__ = ["foo"] 1043 1044class SimpleNewObj(object): 1045 def __init__(self, a, b, c): 1046 # raise an error, to make sure this isn't called 1047 raise TypeError("SimpleNewObj.__init__() didn't expect to get called") 1048 1049class AbstractPickleModuleTests(unittest.TestCase): 1050 1051 def test_dump_closed_file(self): 1052 import os 1053 f = open(TESTFN, "w") 1054 try: 1055 f.close() 1056 self.assertRaises(ValueError, self.module.dump, 123, f) 1057 finally: 1058 os.remove(TESTFN) 1059 1060 def test_load_closed_file(self): 1061 import os 1062 f = open(TESTFN, "w") 1063 try: 1064 f.close() 1065 self.assertRaises(ValueError, self.module.dump, 123, f) 1066 finally: 1067 os.remove(TESTFN) 1068 1069 def test_load_from_and_dump_to_file(self): 1070 stream = cStringIO.StringIO() 1071 data = [123, {}, 124] 1072 self.module.dump(data, stream) 1073 stream.seek(0) 1074 unpickled = self.module.load(stream) 1075 self.assertEqual(unpickled, data) 1076 1077 def test_highest_protocol(self): 1078 # Of course this needs to be changed when HIGHEST_PROTOCOL changes. 1079 self.assertEqual(self.module.HIGHEST_PROTOCOL, 2) 1080 1081 def test_callapi(self): 1082 f = cStringIO.StringIO() 1083 # With and without keyword arguments 1084 self.module.dump(123, f, -1) 1085 self.module.dump(123, file=f, protocol=-1) 1086 self.module.dumps(123, -1) 1087 self.module.dumps(123, protocol=-1) 1088 self.module.Pickler(f, -1) 1089 self.module.Pickler(f, protocol=-1) 1090 1091 def test_incomplete_input(self): 1092 s = cStringIO.StringIO("X''.") 1093 self.assertRaises(EOFError, self.module.load, s) 1094 1095class AbstractPersistentPicklerTests(unittest.TestCase): 1096 1097 # This class defines persistent_id() and persistent_load() 1098 # functions that should be used by the pickler. All even integers 1099 # are pickled using persistent ids. 1100 1101 def persistent_id(self, object): 1102 if isinstance(object, int) and object % 2 == 0: 1103 self.id_count += 1 1104 return str(object) 1105 else: 1106 return None 1107 1108 def persistent_load(self, oid): 1109 self.load_count += 1 1110 object = int(oid) 1111 assert object % 2 == 0 1112 return object 1113 1114 def test_persistence(self): 1115 self.id_count = 0 1116 self.load_count = 0 1117 L = range(10) 1118 self.assertEqual(self.loads(self.dumps(L)), L) 1119 self.assertEqual(self.id_count, 5) 1120 self.assertEqual(self.load_count, 5) 1121 1122 def test_bin_persistence(self): 1123 self.id_count = 0 1124 self.load_count = 0 1125 L = range(10) 1126 self.assertEqual(self.loads(self.dumps(L, 1)), L) 1127 self.assertEqual(self.id_count, 5) 1128 self.assertEqual(self.load_count, 5) 1129 1130class AbstractPicklerUnpicklerObjectTests(unittest.TestCase): 1131 1132 pickler_class = None 1133 unpickler_class = None 1134 1135 def setUp(self): 1136 assert self.pickler_class 1137 assert self.unpickler_class 1138 1139 def dumps(self, obj, proto=0): 1140 buf = cStringIO.StringIO() 1141 pickler = self.pickler_class(buf, proto) 1142 pickler.dump(obj) 1143 return buf.getvalue() 1144 1145 def loads(self, input): 1146 unpickler = self.unpickler_class(cStringIO.StringIO(input)) 1147 return unpickler.load() 1148 1149 def test_clear_pickler_memo(self): 1150 # To test whether clear_memo() has any effect, we pickle an object, 1151 # then pickle it again without clearing the memo; the two serialized 1152 # forms should be different. If we clear_memo() and then pickle the 1153 # object again, the third serialized form should be identical to the 1154 # first one we obtained. 1155 data = ["abcdefg", "abcdefg", 44] 1156 f = cStringIO.StringIO() 1157 pickler = self.pickler_class(f) 1158 1159 pickler.dump(data) 1160 first_pickled = f.getvalue() 1161 1162 # Reset StringIO object. 1163 f.seek(0) 1164 f.truncate() 1165 1166 pickler.dump(data) 1167 second_pickled = f.getvalue() 1168 1169 # Reset the Pickler and StringIO objects. 1170 pickler.clear_memo() 1171 f.seek(0) 1172 f.truncate() 1173 1174 pickler.dump(data) 1175 third_pickled = f.getvalue() 1176 1177 self.assertNotEqual(first_pickled, second_pickled) 1178 self.assertEqual(first_pickled, third_pickled) 1179 1180 def test_priming_pickler_memo(self): 1181 # Verify that we can set the Pickler's memo attribute. 1182 data = ["abcdefg", "abcdefg", 44] 1183 f = cStringIO.StringIO() 1184 pickler = self.pickler_class(f) 1185 1186 pickler.dump(data) 1187 first_pickled = f.getvalue() 1188 1189 f = cStringIO.StringIO() 1190 primed = self.pickler_class(f) 1191 primed.memo = pickler.memo 1192 1193 primed.dump(data) 1194 primed_pickled = f.getvalue() 1195 1196 self.assertNotEqual(first_pickled, primed_pickled) 1197 1198 def test_priming_unpickler_memo(self): 1199 # Verify that we can set the Unpickler's memo attribute. 1200 data = ["abcdefg", "abcdefg", 44] 1201 f = cStringIO.StringIO() 1202 pickler = self.pickler_class(f) 1203 1204 pickler.dump(data) 1205 first_pickled = f.getvalue() 1206 1207 f = cStringIO.StringIO() 1208 primed = self.pickler_class(f) 1209 primed.memo = pickler.memo 1210 1211 primed.dump(data) 1212 primed_pickled = f.getvalue() 1213 1214 unpickler = self.unpickler_class(cStringIO.StringIO(first_pickled)) 1215 unpickled_data1 = unpickler.load() 1216 1217 self.assertEqual(unpickled_data1, data) 1218 1219 primed = self.unpickler_class(cStringIO.StringIO(primed_pickled)) 1220 primed.memo = unpickler.memo 1221 unpickled_data2 = primed.load() 1222 1223 # TODO(collinwinter): add dedicated tests for memo.clear(). This will 1224 # require implementing stream support for unpickling. 1225 primed.memo.clear() 1226 1227 self.assertEqual(unpickled_data2, data) 1228 self.assertTrue(unpickled_data2 is unpickled_data1) 1229 1230 def test_reusing_unpickler_objects(self): 1231 data1 = ["abcdefg", "abcdefg", 44] 1232 f = cStringIO.StringIO() 1233 pickler = self.pickler_class(f) 1234 pickler.dump(data1) 1235 pickled1 = f.getvalue() 1236 1237 data2 = ["abcdefg", 44, 44] 1238 f = cStringIO.StringIO() 1239 pickler = self.pickler_class(f) 1240 pickler.dump(data2) 1241 pickled2 = f.getvalue() 1242 1243 f = cStringIO.StringIO() 1244 f.write(pickled1) 1245 f.seek(0) 1246 unpickler = self.unpickler_class(f) 1247 self.assertEqual(unpickler.load(), data1) 1248 1249 f.seek(0) 1250 f.truncate() 1251 f.write(pickled2) 1252 f.seek(0) 1253 self.assertEqual(unpickler.load(), data2) 1254 1255 def test_pickle_multiple_objects_to_the_same_stream(self): 1256 # Test that we don't optimize out necessary PUT opcodes from the first 1257 # pickle. Even though *that individual pickle* may not need the memo 1258 # state, subsequent pickles written with the same Pickler object are 1259 # relying on the memo state. 1260 val1 = "this is a string" 1261 data = [(val1, 1), (val1, 2), (val1, 3)] 1262 1263 f = cStringIO.StringIO() 1264 pickler = self.pickler_class(f) 1265 for obj in data: 1266 pickler.dump(obj) 1267 1268 f.seek(0) 1269 1270 unpickler = self.unpickler_class(f) 1271 for obj in data: 1272 got_obj = unpickler.load() 1273 self.assertEqual(obj, got_obj) 1274 1275 def test_copy_pickler_memo(self): 1276 pickler = self.pickler_class(cStringIO.StringIO()) 1277 pickler.dump("abadsfasd") 1278 memo = pickler.memo.copy() 1279 pickler.memo.clear() 1280 self.assertEqual(len(memo), 1) 1281 1282 def test_copy_unpickler_memo(self): 1283 string = "abasdfasdf" 1284 sample_data = self.dumps([string, string]) 1285 unpickler = self.unpickler_class(cStringIO.StringIO(sample_data)) 1286 unpickler.load() 1287 memo = unpickler.memo.copy() 1288 self.assertTrue(string in memo.values(), memo) 1289 unpickler.memo.clear() 1290 self.assertTrue(string in memo.values(), memo) 1291 # Make sure we can call copy() and that the unpickler accepts the 1292 # returned object. 1293 unpickler.memo = memo.copy() 1294 1295 def test_priming_picklers_unpicklers(self): 1296 # This test is distilled from the cvs2svn tool. In order to save space 1297 # in their pickle database, cvs2svn "primes" the pickler/unpickler memos 1298 # with common objects so that these objects are represented by GET 1299 # opcodes in the pickles, rather than the much longer sequences they 1300 # would otherwise generate. In order to make use of the priming memos, 1301 # the memos themselves are pickled to the front of the data stream. 1302 # We simulate this process here. Consider this a large integration test, 1303 # while the tests above are intended to be more fine-grained. 1304 1305 # Objects we want to share between pickles. 1306 common1 = "asdfasdfasdfa" 1307 common2 = "asdfadfadfshdifu" 1308 common3 = (5, 6) 1309 common_objects = [common1, common2, common3] 1310 1311 # Combinations of the above objects. These will be sent through 1312 # the data stream. 1313 combo1 = (common1, common2) 1314 combo2 = (6, common3) 1315 combo3 = (common1, common2, common3) 1316 1317 data_stream = cStringIO.StringIO() 1318 1319 # 1. Pickle the common objects, extract the Pickler's memo. 1320 f = cStringIO.StringIO() 1321 pickler = self.pickler_class(f) 1322 pickler.dump(common_objects) 1323 pickler_memo = pickler.memo.copy() 1324 priming_data = f.getvalue() 1325 1326 # 2. Unpickle the common objects, extract the Unpickler's memo. 1327 unpickler = self.unpickler_class(cStringIO.StringIO(priming_data)) 1328 unpickler.load() 1329 unpickler_memo = unpickler.memo.copy() 1330 1331 # 3. Write the pickler and unpickler memos to the data stream. These 1332 # will be used to retrieve data from the stream. 1333 pickler = self.pickler_class(data_stream) 1334 pickler.dump(pickler_memo) 1335 pickler.dump(unpickler_memo) 1336 1337 # 4. Pickle some more data to the data stream, using the priming memos. 1338 # cvs2svn resets the memo before every object. 1339 pickler = self.pickler_class(data_stream) 1340 pickler.memo = pickler_memo.copy() 1341 pickler.dump(combo1) 1342 pickler.memo = pickler_memo.copy() 1343 pickler.dump(combo2) 1344 1345 # 5. Now start reading the data. We read the priming memos first. 1346 data_stream.seek(0) 1347 unpickler = self.unpickler_class(data_stream) 1348 recv_pickler_memo = unpickler.load() 1349 recv_unpickler_memo = unpickler.load() 1350 self.assertEqual(recv_pickler_memo, pickler_memo) 1351 self.assertEqual(recv_unpickler_memo, unpickler_memo) 1352 1353 # 6. Using the received unpickling primer memo, unpickle the rest of 1354 # the data stream. cvs2svn resets the memo before every object. 1355 unpickler = self.unpickler_class(data_stream) 1356 unpickler.memo = recv_unpickler_memo.copy() 1357 self.assertEqual(unpickler.load(), combo1) 1358 unpickler.memo = recv_unpickler_memo.copy() 1359 self.assertEqual(unpickler.load(), combo2) 1360 1361 # 7. Pickle something using the received pickler memo, then unpickle it 1362 # to make sure that the pickler memo works. 1363 f = cStringIO.StringIO() 1364 pickler = self.pickler_class(f) 1365 pickler.memo = recv_pickler_memo.copy() 1366 pickler.dump(combo3) 1367 f.seek(0) 1368 unpickler = self.unpickler_class(f) 1369 unpickler.memo = recv_unpickler_memo.copy() 1370 self.assertEqual(unpickler.load(), combo3)