Lib/test/pickletester.py

   1 import unittest
   2 import pickle
   3 import cPickle
   4 import pickletools
   5 import copyreg
   6
   7 from test.test_support import TestFailed, have_unicode, TESTFN, \
   8                               run_with_locale
   9
  10 # Tests that try a number of pickle protocols should have a
  11 #     for proto in protocols:
  12 # kind of outer loop.
  13 assert pickle.HIGHEST_PROTOCOL == cPickle.HIGHEST_PROTOCOL == 2
  14 protocols = range(pickle.HIGHEST_PROTOCOL + 1)
  15
  16
  17 # Return True if opcode code appears in the pickle, else False.
  18 def opcode_in_pickle(code, pickle):
  19     for op, dummy, dummy in pickletools.genops(pickle):
  20         if op.code == code:
  21             return True
  22     return False
  23
  24 # Return the number of times opcode code appears in pickle.
  25 def count_opcode(code, pickle):
  26     n = 0
  27     for op, dummy, dummy in pickletools.genops(pickle):
  28         if op.code == code:
  29             n += 1
  30     return n
  31
  32 # We can't very well test the extension registry without putting known stuff
  33 # in it, but we have to be careful to restore its original state.  Code
  34 # should do this:
  35 #
  36 #     e = ExtensionSaver(extension_code)
  37 #     try:
  38 #         fiddle w/ the extension registry's stuff for extension_code
  39 #     finally:
  40 #         e.restore()
  41
  42 class ExtensionSaver:
  43     # Remember current registration for code (if any), and remove it (if
  44     # there is one).
  45     def __init__(self, code):
  46         self.code = code
  47         if code in copyreg._inverted_registry:
  48             self.pair = copyreg._inverted_registry[code]
  49             copyreg.remove_extension(self.pair[0], self.pair[1], code)
  50         else:
  51             self.pair = None
  52
  53     # Restore previous registration for code.
  54     def restore(self):
  55         code = self.code
  56         curpair = copyreg._inverted_registry.get(code)
  57         if curpair is not None:
  58             copyreg.remove_extension(curpair[0], curpair[1], code)
  59         pair = self.pair
  60         if pair is not None:
  61             copyreg.add_extension(pair[0], pair[1], code)
  62
  63 class C:
  64     def __cmp__(self, other):
  65         return cmp(self.__dict__, other.__dict__)
  66
  67 import __main__
  68 __main__.C = C
  69 C.__module__ = "__main__"
  70
  71 class myint(int):
  72     def __init__(self, x):
  73         self.str = str(x)
  74
  75 class initarg(C):
  76
  77     def __init__(self, a, b):
  78         self.a = a
  79         self.b = b
  80
  81     def __getinitargs__(self):
  82         return self.a, self.b
  83
  84 class metaclass(type):
  85     pass
  86
  87 class use_metaclass(object):
  88     __metaclass__ = metaclass
  89
  90 # DATA0 .. DATA2 are the pickles we expect under the various protocols, for
  91 # the object returned by create_data().
  92
  93 # break into multiple strings to avoid confusing font-lock-mode
  94 DATA0 = """(lp1
  95 I0
  96 aL1L
  97 aF2
  98 ac__builtin__
  99 complex
 100 p2
 101 """ + \
 102 """(F3
 103 F0
 104 tRp3
 105 aI1
 106 aI-1
 107 aI255
 108 aI-255
 109 aI-256
 110 aI65535
 111 aI-65535
 112 aI-65536
 113 aI2147483647
 114 aI-2147483647
 115 aI-2147483648
 116 a""" + \
 117 """(S'abc'
 118 p4
 119 g4
 120 """ + \
 121 """(i__main__
 122 C
 123 p5
 124 """ + \
 125 """(dp6
 126 S'foo'
 127 p7
 128 I1
 129 sS'bar'
 130 p8
 131 I2
 132 sbg5
 133 tp9
 134 ag9
 135 aI5
 136 a.
 137 """
 138
 139 # Disassembly of DATA0.
 140 DATA0_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
 209 highest protocol among opcodes = 0
 210 """
 211
 212 DATA1 = (']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.
 223 DATA1_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
 276 highest protocol among opcodes = 1
 277 """
 278
 279 DATA2 = ('\x80\x02]q\x01(K\x00\x8a\x01\x01G@\x00\x00\x00\x00\x00\x00\x00'
 280          'c__builtin__\ncomplex\nq\x02G@\x08\x00\x00\x00\x00\x00\x00G\x00'
 281          '\x00\x00\x00\x00\x00\x00\x00\x86Rq\x03K\x01J\xff\xff\xff\xffK'
 282          '\xffJ\x01\xff\xff\xffJ\x00\xff\xff\xffM\xff\xffJ\x01\x00\xff\xff'
 283          'J\x00\x00\xff\xffJ\xff\xff\xff\x7fJ\x01\x00\x00\x80J\x00\x00\x00'
 284          '\x80(U\x03abcq\x04h\x04(c__main__\nC\nq\x05oq\x06}q\x07(U\x03foo'
 285          'q\x08K\x01U\x03barq\tK\x02ubh\x06tq\nh\nK\x05e.')
 286
 287 # Disassembly of DATA2.
 288 DATA2_DIS = """\
 289     0: \x80 PROTO      2
 290     2: ]    EMPTY_LIST
 291     3: q    BINPUT     1
 292     5: (    MARK
 293     6: K        BININT1    0
 294     8: \x8a     LONG1      1L
 295    11: G        BINFLOAT   2.0
 296    20: c        GLOBAL     '__builtin__ complex'
 297    41: q        BINPUT     2
 298    43: G        BINFLOAT   3.0
 299    52: G        BINFLOAT   0.0
 300    61: \x86     TUPLE2
 301    62: R        REDUCE
 302    63: q        BINPUT     3
 303    65: K        BININT1    1
 304    67: J        BININT     -1
 305    72: K        BININT1    255
 306    74: J        BININT     -255
 307    79: J        BININT     -256
 308    84: M        BININT2    65535
 309    87: J        BININT     -65535
 310    92: J        BININT     -65536
 311    97: J        BININT     2147483647
 312   102: J        BININT     -2147483647
 313   107: J        BININT     -2147483648
 314   112: (        MARK
 315   113: U            SHORT_BINSTRING 'abc'
 316   118: q            BINPUT     4
 317   120: h            BINGET     4
 318   122: (            MARK
 319   123: c                GLOBAL     '__main__ C'
 320   135: q                BINPUT     5
 321   137: o                OBJ        (MARK at 122)
 322   138: q            BINPUT     6
 323   140: }            EMPTY_DICT
 324   141: q            BINPUT     7
 325   143: (            MARK
 326   144: U                SHORT_BINSTRING 'foo'
 327   149: q                BINPUT     8
 328   151: K                BININT1    1
 329   153: U                SHORT_BINSTRING 'bar'
 330   158: q                BINPUT     9
 331   160: K                BININT1    2
 332   162: u                SETITEMS   (MARK at 143)
 333   163: b            BUILD
 334   164: h            BINGET     6
 335   166: t            TUPLE      (MARK at 112)
 336   167: q        BINPUT     10
 337   169: h        BINGET     10
 338   171: K        BININT1    5
 339   173: e        APPENDS    (MARK at 5)
 340   174: .    STOP
 341 highest protocol among opcodes = 2
 342 """
 343
 344 def create_data():
 345     c = C()
 346     c.foo = 1
 347     c.bar = 2
 348     x = [0, 1L, 2.0, 3.0+0j]
 349     # Append some integer test cases at cPickle.c's internal size
 350     # cutoffs.
 351     uint1max = 0xff
 352     uint2max = 0xffff
 353     int4max = 0x7fffffff
 354     x.extend([1, -1,
 355               uint1max, -uint1max, -uint1max-1,
 356               uint2max, -uint2max, -uint2max-1,
 357                int4max,  -int4max,  -int4max-1])
 358     y = ('abc', 'abc', c, c)
 359     x.append(y)
 360     x.append(y)
 361     x.append(5)
 362     return x
 363
 364 class AbstractPickleTests(unittest.TestCase):
 365     # Subclass must define self.dumps, self.loads, self.error.
 366
 367     _testdata = create_data()
 368
 369     def setUp(self):
 370         pass
 371
 372     def test_misc(self):
 373         # test various datatypes not tested by testdata
 374         for proto in protocols:
 375             x = myint(4)
 376             s = self.dumps(x, proto)
 377             y = self.loads(s)
 378             self.assertEqual(x, y)
 379
 380             x = (1, ())
 381             s = self.dumps(x, proto)
 382             y = self.loads(s)
 383             self.assertEqual(x, y)
 384
 385             x = initarg(1, x)
 386             s = self.dumps(x, proto)
 387             y = self.loads(s)
 388             self.assertEqual(x, y)
 389
 390         # XXX test __reduce__ protocol?
 391
 392     def test_roundtrip_equality(self):
 393         expected = self._testdata
 394         for proto in protocols:
 395             s = self.dumps(expected, proto)
 396             got = self.loads(s)
 397             self.assertEqual(expected, got)
 398
 399     def test_load_from_canned_string(self):
 400         expected = self._testdata
 401         for canned in DATA0, DATA1, DATA2:
 402             got = self.loads(canned)
 403             self.assertEqual(expected, got)
 404
 405     # There are gratuitous differences between pickles produced by
 406     # pickle and cPickle, largely because cPickle starts PUT indices at
 407     # 1 and pickle starts them at 0.  See XXX comment in cPickle's put2() --
 408     # there's a comment with an exclamation point there whose meaning
 409     # is a mystery.  cPickle also suppresses PUT for objects with a refcount
 410     # of 1.
 411     def dont_test_disassembly(self):
 412         from cStringIO import StringIO
 413         from pickletools import dis
 414
 415         for proto, expected in (0, DATA0_DIS), (1, DATA1_DIS):
 416             s = self.dumps(self._testdata, proto)
 417             filelike = StringIO()
 418             dis(s, out=filelike)
 419             got = filelike.getvalue()
 420             self.assertEqual(expected, got)
 421
 422     def test_recursive_list(self):
 423         l = []
 424         l.append(l)
 425         for proto in protocols:
 426             s = self.dumps(l, proto)
 427             x = self.loads(s)
 428             self.assertEqual(len(x), 1)
 429             self.assert_(x is x[0])
 430
 431     def test_recursive_dict(self):
 432         d = {}
 433         d[1] = d
 434         for proto in protocols:
 435             s = self.dumps(d, proto)
 436             x = self.loads(s)
 437             self.assertEqual(x.keys(), [1])
 438             self.assert_(x[1] is x)
 439
 440     def test_recursive_inst(self):
 441         i = C()
 442         i.attr = i
 443         for proto in protocols:
 444             s = self.dumps(i, 2)
 445             x = self.loads(s)
 446             self.assertEqual(dir(x), dir(i))
 447             self.assert_(x.attr is x)
 448
 449     def test_recursive_multi(self):
 450         l = []
 451         d = {1:l}
 452         i = C()
 453         i.attr = d
 454         l.append(i)
 455         for proto in protocols:
 456             s = self.dumps(l, proto)
 457             x = self.loads(s)
 458             self.assertEqual(len(x), 1)
 459             self.assertEqual(dir(x[0]), dir(i))
 460             self.assertEqual(x[0].attr.keys(), [1])
 461             self.assert_(x[0].attr[1] is x)
 462
 463     def test_garyp(self):
 464         self.assertRaises(self.error, self.loads, 'garyp')
 465
 466     def test_insecure_strings(self):
 467         insecure = ["abc", "2 + 2", # not quoted
 468                     #"'abc' + 'def'", # not a single quoted string
 469                     "'abc", # quote is not closed
 470                     "'abc\"", # open quote and close quote don't match
 471                     "'abc'   ?", # junk after close quote
 472                     "'\\'", # trailing backslash
 473                     # some tests of the quoting rules
 474                     #"'abc\"\''",
 475                     #"'\\\\a\'\'\'\\\'\\\\\''",
 476                     ]
 477         for s in insecure:
 478             buf = "S" + s + "\012p0\012."
 479             self.assertRaises(ValueError, self.loads, buf)
 480
 481     if have_unicode:
 482         def test_unicode(self):
 483             endcases = [unicode(''), unicode('<\\u>'), unicode('<\\\u1234>'),
 484                         unicode('<\n>'),  unicode('<\\>')]
 485             for proto in protocols:
 486                 for u in endcases:
 487                     p = self.dumps(u, proto)
 488                     u2 = self.loads(p)
 489                     self.assertEqual(u2, u)
 490
 491     def test_ints(self):
 492         import sys
 493         for proto in protocols:
 494             n = sys.maxint
 495             while n:
 496                 for expected in (-n, n):
 497                     s = self.dumps(expected, proto)
 498                     n2 = self.loads(s)
 499                     self.assertEqual(expected, n2)
 500                 n = n >> 1
 501
 502     def test_maxint64(self):
 503         maxint64 = (1L << 63) - 1
 504         data = 'I' + str(maxint64) + '\n.'
 505         got = self.loads(data)
 506         self.assertEqual(got, maxint64)
 507
 508         # Try too with a bogus literal.
 509         data = 'I' + str(maxint64) + 'JUNK\n.'
 510         self.assertRaises(ValueError, self.loads, data)
 511
 512     def test_long(self):
 513         for proto in protocols:
 514             # 256 bytes is where LONG4 begins.
 515             for nbits in 1, 8, 8*254, 8*255, 8*256, 8*257:
 516                 nbase = 1L << nbits
 517                 for npos in nbase-1, nbase, nbase+1:
 518                     for n in npos, -npos:
 519                         pickle = self.dumps(n, proto)
 520                         got = self.loads(pickle)
 521                         self.assertEqual(n, got)
 522         # Try a monster.  This is quadratic-time in protos 0 & 1, so don't
 523         # bother with those.
 524         nbase = long("deadbeeffeedface", 16)
 525         nbase += nbase << 1000000
 526         for n in nbase, -nbase:
 527             p = self.dumps(n, 2)
 528             got = self.loads(p)
 529             self.assertEqual(n, got)
 530
 531     @run_with_locale('LC_ALL', 'de_DE', 'fr_FR')
 532     def test_float_format(self):
 533         # make sure that floats are formatted locale independent
 534         self.assertEqual(self.dumps(1.2)[0:3], 'F1.')
 535
 536     def test_reduce(self):
 537         pass
 538
 539     def test_getinitargs(self):
 540         pass
 541
 542     def test_metaclass(self):
 543         a = use_metaclass()
 544         for proto in protocols:
 545             s = self.dumps(a, proto)
 546             b = self.loads(s)
 547             self.assertEqual(a.__class__, b.__class__)
 548
 549     def test_structseq(self):
 550         import time
 551         import os
 552
 553         t = time.localtime()
 554         for proto in protocols:
 555             s = self.dumps(t, proto)
 556             u = self.loads(s)
 557             self.assertEqual(t, u)
 558             if hasattr(os, "stat"):
 559                 t = os.stat(os.curdir)
 560                 s = self.dumps(t, proto)
 561                 u = self.loads(s)
 562                 self.assertEqual(t, u)
 563             if hasattr(os, "statvfs"):
 564                 t = os.statvfs(os.curdir)
 565                 s = self.dumps(t, proto)
 566                 u = self.loads(s)
 567                 self.assertEqual(t, u)
 568
 569     # Tests for protocol 2
 570
 571     def test_proto(self):
 572         build_none = pickle.NONE + pickle.STOP
 573         for proto in protocols:
 574             expected = build_none
 575             if proto >= 2:
 576                 expected = pickle.PROTO + chr(proto) + expected
 577             p = self.dumps(None, proto)
 578             self.assertEqual(p, expected)
 579
 580         oob = protocols[-1] + 1     # a future protocol
 581         badpickle = pickle.PROTO + chr(oob) + build_none
 582         try:
 583             self.loads(badpickle)
 584         except ValueError, detail:
 585             self.failUnless(str(detail).startswith(
 586                                             "unsupported pickle protocol"))
 587         else:
 588             self.fail("expected bad protocol number to raise ValueError")
 589
 590     def test_long1(self):
 591         x = 12345678910111213141516178920L
 592         for proto in protocols:
 593             s = self.dumps(x, proto)
 594             y = self.loads(s)
 595             self.assertEqual(x, y)
 596             self.assertEqual(opcode_in_pickle(pickle.LONG1, s), proto >= 2)
 597
 598     def test_long4(self):
 599         x = 12345678910111213141516178920L << (256*8)
 600         for proto in protocols:
 601             s = self.dumps(x, proto)
 602             y = self.loads(s)
 603             self.assertEqual(x, y)
 604             self.assertEqual(opcode_in_pickle(pickle.LONG4, s), proto >= 2)
 605
 606     def test_short_tuples(self):
 607         # Map (proto, len(tuple)) to expected opcode.
 608         expected_opcode = {(0, 0): pickle.TUPLE,
 609                            (0, 1): pickle.TUPLE,
 610                            (0, 2): pickle.TUPLE,
 611                            (0, 3): pickle.TUPLE,
 612                            (0, 4): pickle.TUPLE,
 613
 614                            (1, 0): pickle.EMPTY_TUPLE,
 615                            (1, 1): pickle.TUPLE,
 616                            (1, 2): pickle.TUPLE,
 617                            (1, 3): pickle.TUPLE,
 618                            (1, 4): pickle.TUPLE,
 619
 620                            (2, 0): pickle.EMPTY_TUPLE,
 621                            (2, 1): pickle.TUPLE1,
 622                            (2, 2): pickle.TUPLE2,
 623                            (2, 3): pickle.TUPLE3,
 624                            (2, 4): pickle.TUPLE,
 625                           }
 626         a = ()
 627         b = (1,)
 628         c = (1, 2)
 629         d = (1, 2, 3)
 630         e = (1, 2, 3, 4)
 631         for proto in protocols:
 632             for x in a, b, c, d, e:
 633                 s = self.dumps(x, proto)
 634                 y = self.loads(s)
 635                 self.assertEqual(x, y, (proto, x, s, y))
 636                 expected = expected_opcode[proto, len(x)]
 637                 self.assertEqual(opcode_in_pickle(expected, s), True)
 638
 639     def test_singletons(self):
 640         # Map (proto, singleton) to expected opcode.
 641         expected_opcode = {(0, None): pickle.NONE,
 642                            (1, None): pickle.NONE,
 643                            (2, None): pickle.NONE,
 644
 645                            (0, True): pickle.INT,
 646                            (1, True): pickle.INT,
 647                            (2, True): pickle.NEWTRUE,
 648
 649                            (0, False): pickle.INT,
 650                            (1, False): pickle.INT,
 651                            (2, False): pickle.NEWFALSE,
 652                           }
 653         for proto in protocols:
 654             for x in None, False, True:
 655                 s = self.dumps(x, proto)
 656                 y = self.loads(s)
 657                 self.assert_(x is y, (proto, x, s, y))
 658                 expected = expected_opcode[proto, x]
 659                 self.assertEqual(opcode_in_pickle(expected, s), True)
 660
 661     def test_newobj_tuple(self):
 662         x = MyTuple([1, 2, 3])
 663         x.foo = 42
 664         x.bar = "hello"
 665         for proto in protocols:
 666             s = self.dumps(x, proto)
 667             y = self.loads(s)
 668             self.assertEqual(tuple(x), tuple(y))
 669             self.assertEqual(x.__dict__, y.__dict__)
 670
 671     def test_newobj_list(self):
 672         x = MyList([1, 2, 3])
 673         x.foo = 42
 674         x.bar = "hello"
 675         for proto in protocols:
 676             s = self.dumps(x, proto)
 677             y = self.loads(s)
 678             self.assertEqual(list(x), list(y))
 679             self.assertEqual(x.__dict__, y.__dict__)
 680
 681     def test_newobj_generic(self):
 682         for proto in protocols:
 683             for C in myclasses:
 684                 B = C.__base__
 685                 x = C(C.sample)
 686                 x.foo = 42
 687                 s = self.dumps(x, proto)
 688                 y = self.loads(s)
 689                 detail = (proto, C, B, x, y, type(y))
 690                 self.assertEqual(B(x), B(y), detail)
 691                 self.assertEqual(x.__dict__, y.__dict__, detail)
 692
 693     # Register a type with copyreg, with extension code extcode.  Pickle
 694     # an object of that type.  Check that the resulting pickle uses opcode
 695     # (EXT[124]) under proto 2, and not in proto 1.
 696
 697     def produce_global_ext(self, extcode, opcode):
 698         e = ExtensionSaver(extcode)
 699         try:
 700             copyreg.add_extension(__name__, "MyList", extcode)
 701             x = MyList([1, 2, 3])
 702             x.foo = 42
 703             x.bar = "hello"
 704
 705             # Dump using protocol 1 for comparison.
 706             s1 = self.dumps(x, 1)
 707             self.assert_(__name__ in s1)
 708             self.assert_("MyList" in s1)
 709             self.assertEqual(opcode_in_pickle(opcode, s1), False)
 710
 711             y = self.loads(s1)
 712             self.assertEqual(list(x), list(y))
 713             self.assertEqual(x.__dict__, y.__dict__)
 714
 715             # Dump using protocol 2 for test.
 716             s2 = self.dumps(x, 2)
 717             self.assert_(__name__ not in s2)
 718             self.assert_("MyList" not in s2)
 719             self.assertEqual(opcode_in_pickle(opcode, s2), True)
 720
 721             y = self.loads(s2)
 722             self.assertEqual(list(x), list(y))
 723             self.assertEqual(x.__dict__, y.__dict__)
 724
 725         finally:
 726             e.restore()
 727
 728     def test_global_ext1(self):
 729         self.produce_global_ext(0x00000001, pickle.EXT1)  # smallest EXT1 code
 730         self.produce_global_ext(0x000000ff, pickle.EXT1)  # largest EXT1 code
 731
 732     def test_global_ext2(self):
 733         self.produce_global_ext(0x00000100, pickle.EXT2)  # smallest EXT2 code
 734         self.produce_global_ext(0x0000ffff, pickle.EXT2)  # largest EXT2 code
 735         self.produce_global_ext(0x0000abcd, pickle.EXT2)  # check endianness
 736
 737     def test_global_ext4(self):
 738         self.produce_global_ext(0x00010000, pickle.EXT4)  # smallest EXT4 code
 739         self.produce_global_ext(0x7fffffff, pickle.EXT4)  # largest EXT4 code
 740         self.produce_global_ext(0x12abcdef, pickle.EXT4)  # check endianness
 741
 742     def test_list_chunking(self):
 743         n = 10  # too small to chunk
 744         x = range(n)
 745         for proto in protocols:
 746             s = self.dumps(x, proto)
 747             y = self.loads(s)
 748             self.assertEqual(x, y)
 749             num_appends = count_opcode(pickle.APPENDS, s)
 750             self.assertEqual(num_appends, proto > 0)
 751
 752         n = 2500  # expect at least two chunks when proto > 0
 753         x = range(n)
 754         for proto in protocols:
 755             s = self.dumps(x, proto)
 756             y = self.loads(s)
 757             self.assertEqual(x, y)
 758             num_appends = count_opcode(pickle.APPENDS, s)
 759             if proto == 0:
 760                 self.assertEqual(num_appends, 0)
 761             else:
 762                 self.failUnless(num_appends >= 2)
 763
 764     def test_dict_chunking(self):
 765         n = 10  # too small to chunk
 766         x = dict.fromkeys(range(n))
 767         for proto in protocols:
 768             s = self.dumps(x, proto)
 769             y = self.loads(s)
 770             self.assertEqual(x, y)
 771             num_setitems = count_opcode(pickle.SETITEMS, s)
 772             self.assertEqual(num_setitems, proto > 0)
 773
 774         n = 2500  # expect at least two chunks when proto > 0
 775         x = dict.fromkeys(range(n))
 776         for proto in protocols:
 777             s = self.dumps(x, proto)
 778             y = self.loads(s)
 779             self.assertEqual(x, y)
 780             num_setitems = count_opcode(pickle.SETITEMS, s)
 781             if proto == 0:
 782                 self.assertEqual(num_setitems, 0)
 783             else:
 784                 self.failUnless(num_setitems >= 2)
 785
 786     def test_simple_newobj(self):
 787         x = object.__new__(SimpleNewObj)  # avoid __init__
 788         x.abc = 666
 789         for proto in protocols:
 790             s = self.dumps(x, proto)
 791             self.assertEqual(opcode_in_pickle(pickle.NEWOBJ, s), proto >= 2)
 792             y = self.loads(s)   # will raise TypeError if __init__ called
 793             self.assertEqual(y.abc, 666)
 794             self.assertEqual(x.__dict__, y.__dict__)
 795
 796     def test_newobj_list_slots(self):
 797         x = SlotList([1, 2, 3])
 798         x.foo = 42
 799         x.bar = "hello"
 800         s = self.dumps(x, 2)
 801         y = self.loads(s)
 802         self.assertEqual(list(x), list(y))
 803         self.assertEqual(x.__dict__, y.__dict__)
 804         self.assertEqual(x.foo, y.foo)
 805         self.assertEqual(x.bar, y.bar)
 806
 807     def test_reduce_overrides_default_reduce_ex(self):
 808         for proto in 0, 1, 2:
 809             x = REX_one()
 810             self.assertEqual(x._reduce_called, 0)
 811             s = self.dumps(x, proto)
 812             self.assertEqual(x._reduce_called, 1)
 813             y = self.loads(s)
 814             self.assertEqual(y._reduce_called, 0)
 815
 816     def test_reduce_ex_called(self):
 817         for proto in 0, 1, 2:
 818             x = REX_two()
 819             self.assertEqual(x._proto, None)
 820             s = self.dumps(x, proto)
 821             self.assertEqual(x._proto, proto)
 822             y = self.loads(s)
 823             self.assertEqual(y._proto, None)
 824
 825     def test_reduce_ex_overrides_reduce(self):
 826         for proto in 0, 1, 2:
 827             x = REX_three()
 828             self.assertEqual(x._proto, None)
 829             s = self.dumps(x, proto)
 830             self.assertEqual(x._proto, proto)
 831             y = self.loads(s)
 832             self.assertEqual(y._proto, None)
 833
 834     def test_reduce_ex_calls_base(self):
 835         for proto in 0, 1, 2:
 836             x = REX_four()
 837             self.assertEqual(x._proto, None)
 838             s = self.dumps(x, proto)
 839             self.assertEqual(x._proto, proto)
 840             y = self.loads(s)
 841             self.assertEqual(y._proto, proto)
 842
 843     def test_reduce_calls_base(self):
 844         for proto in 0, 1, 2:
 845             x = REX_five()
 846             self.assertEqual(x._reduce_called, 0)
 847             s = self.dumps(x, proto)
 848             self.assertEqual(x._reduce_called, 1)
 849             y = self.loads(s)
 850             self.assertEqual(y._reduce_called, 1)
 851
 852 # Test classes for reduce_ex
 853
 854 class REX_one(object):
 855     _reduce_called = 0
 856     def __reduce__(self):
 857         self._reduce_called = 1
 858         return REX_one, ()
 859     # No __reduce_ex__ here, but inheriting it from object
 860
 861 class REX_two(object):
 862     _proto = None
 863     def __reduce_ex__(self, proto):
 864         self._proto = proto
 865         return REX_two, ()
 866     # No __reduce__ here, but inheriting it from object
 867
 868 class REX_three(object):
 869     _proto = None
 870     def __reduce_ex__(self, proto):
 871         self._proto = proto
 872         return REX_two, ()
 873     def __reduce__(self):
 874         raise TestFailed, "This __reduce__ shouldn't be called"
 875
 876 class REX_four(object):
 877     _proto = None
 878     def __reduce_ex__(self, proto):
 879         self._proto = proto
 880         return object.__reduce_ex__(self, proto)
 881     # Calling base class method should succeed
 882
 883 class REX_five(object):
 884     _reduce_called = 0
 885     def __reduce__(self):
 886         self._reduce_called = 1
 887         return object.__reduce__(self)
 888     # This one used to fail with infinite recursion
 889
 890 # Test classes for newobj
 891
 892 class MyInt(int):
 893     sample = 1
 894
 895 class MyLong(long):
 896     sample = 1L
 897
 898 class MyFloat(float):
 899     sample = 1.0
 900
 901 class MyComplex(complex):
 902     sample = 1.0 + 0.0j
 903
 904 class MyStr(str):
 905     sample = "hello"
 906
 907 class MyUnicode(unicode):
 908     sample = u"hello \u1234"
 909
 910 class MyTuple(tuple):
 911     sample = (1, 2, 3)
 912
 913 class MyList(list):
 914     sample = [1, 2, 3]
 915
 916 class MyDict(dict):
 917     sample = {"a": 1, "b": 2}
 918
 919 myclasses = [MyInt, MyLong, MyFloat,
 920              MyComplex,
 921              MyStr, MyUnicode,
 922              MyTuple, MyList, MyDict]
 923
 924
 925 class SlotList(MyList):
 926     __slots__ = ["foo"]
 927
 928 class SimpleNewObj(object):
 929     def __init__(self, a, b, c):
 930         # raise an error, to make sure this isn't called
 931         raise TypeError("SimpleNewObj.__init__() didn't expect to get called")
 932
 933 class AbstractPickleModuleTests(unittest.TestCase):
 934
 935     def test_dump_closed_file(self):
 936         import os
 937         f = open(TESTFN, "w")
 938         try:
 939             f.close()
 940             self.assertRaises(ValueError, self.module.dump, 123, f)
 941         finally:
 942             os.remove(TESTFN)
 943
 944     def test_load_closed_file(self):
 945         import os
 946         f = open(TESTFN, "w")
 947         try:
 948             f.close()
 949             self.assertRaises(ValueError, self.module.dump, 123, f)
 950         finally:
 951             os.remove(TESTFN)
 952
 953     def test_highest_protocol(self):
 954         # Of course this needs to be changed when HIGHEST_PROTOCOL changes.
 955         self.assertEqual(self.module.HIGHEST_PROTOCOL, 2)
 956
 957     def test_callapi(self):
 958         from cStringIO import StringIO
 959         f = StringIO()
 960         # With and without keyword arguments
 961         self.module.dump(123, f, -1)
 962         self.module.dump(123, file=f, protocol=-1)
 963         self.module.dumps(123, -1)
 964         self.module.dumps(123, protocol=-1)
 965         self.module.Pickler(f, -1)
 966         self.module.Pickler(f, protocol=-1)
 967
 968 class AbstractPersistentPicklerTests(unittest.TestCase):
 969
 970     # This class defines persistent_id() and persistent_load()
 971     # functions that should be used by the pickler.  All even integers
 972     # are pickled using persistent ids.
 973
 974     def persistent_id(self, object):
 975         if isinstance(object, int) and object % 2 == 0:
 976             self.id_count += 1
 977             return str(object)
 978         else:
 979             return None
 980
 981     def persistent_load(self, oid):
 982         self.load_count += 1
 983         object = int(oid)
 984         assert object % 2 == 0
 985         return object
 986
 987     def test_persistence(self):
 988         self.id_count = 0
 989         self.load_count = 0
 990         L = range(10)
 991         self.assertEqual(self.loads(self.dumps(L)), L)
 992         self.assertEqual(self.id_count, 5)
 993         self.assertEqual(self.load_count, 5)
 994
 995     def test_bin_persistence(self):
 996         self.id_count = 0
 997         self.load_count = 0
 998         L = range(10)
 999         self.assertEqual(self.loads(self.dumps(L, 1)), L)
1000         self.assertEqual(self.id_count, 5)
1001         self.assertEqual(self.load_count, 5)