Lib/test/pickletester.py

   1 import unittest
   2 import pickle
   3 import cPickle
   4 import pickletools
   5 import copy_reg
   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 copy_reg._inverted_registry:
  48             self.pair = copy_reg._inverted_registry[code]
  49             copy_reg.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 = copy_reg._inverted_registry.get(code)
  57         if curpair is not None:
  58             copy_reg.remove_extension(curpair[0], curpair[1], code)
  59         pair = self.pair
  60         if pair is not None:
  61             copy_reg.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 = [u'', u'<\\u>', u'<\\\u1234>', u'<\n>',
 484                         u'<\\>', u'<\\\U00012345>']
 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_unicode_high_plane(self):
 492             t = u'\U00012345'
 493             for proto in protocols:
 494                 p = self.dumps(t, proto)
 495                 t2 = self.loads(p)
 496                 self.assertEqual(t2, t)
 497
 498     def test_ints(self):
 499         import sys
 500         for proto in protocols:
 501             n = sys.maxint
 502             while n:
 503                 for expected in (-n, n):
 504                     s = self.dumps(expected, proto)
 505                     n2 = self.loads(s)
 506                     self.assertEqual(expected, n2)
 507                 n = n >> 1
 508
 509     def test_maxint64(self):
 510         maxint64 = (1L << 63) - 1
 511         data = 'I' + str(maxint64) + '\n.'
 512         got = self.loads(data)
 513         self.assertEqual(got, maxint64)
 514
 515         # Try too with a bogus literal.
 516         data = 'I' + str(maxint64) + 'JUNK\n.'
 517         self.assertRaises(ValueError, self.loads, data)
 518
 519     def test_long(self):
 520         for proto in protocols:
 521             # 256 bytes is where LONG4 begins.
 522             for nbits in 1, 8, 8*254, 8*255, 8*256, 8*257:
 523                 nbase = 1L << nbits
 524                 for npos in nbase-1, nbase, nbase+1:
 525                     for n in npos, -npos:
 526                         pickle = self.dumps(n, proto)
 527                         got = self.loads(pickle)
 528                         self.assertEqual(n, got)
 529         # Try a monster.  This is quadratic-time in protos 0 & 1, so don't
 530         # bother with those.
 531         nbase = long("deadbeeffeedface", 16)
 532         nbase += nbase << 1000000
 533         for n in nbase, -nbase:
 534             p = self.dumps(n, 2)
 535             got = self.loads(p)
 536             self.assertEqual(n, got)
 537
 538     def test_float(self):
 539         test_values = [0.0, 4.94e-324, 1e-310, 7e-308, 6.626e-34, 0.1, 0.5,
 540                        3.14, 263.44582062374053, 6.022e23, 1e30]
 541         test_values = test_values + [-x for x in test_values]
 542         for proto in protocols:
 543             for value in test_values:
 544                 pickle = self.dumps(value, proto)
 545                 got = self.loads(pickle)
 546                 self.assertEqual(value, got)
 547
 548     @run_with_locale('LC_ALL', 'de_DE', 'fr_FR')
 549     def test_float_format(self):
 550         # make sure that floats are formatted locale independent
 551         self.assertEqual(self.dumps(1.2)[0:3], 'F1.')
 552
 553     def test_reduce(self):
 554         pass
 555
 556     def test_getinitargs(self):
 557         pass
 558
 559     def test_metaclass(self):
 560         a = use_metaclass()
 561         for proto in protocols:
 562             s = self.dumps(a, proto)
 563             b = self.loads(s)
 564             self.assertEqual(a.__class__, b.__class__)
 565
 566     def test_structseq(self):
 567         import time
 568         import os
 569
 570         t = time.localtime()
 571         for proto in protocols:
 572             s = self.dumps(t, proto)
 573             u = self.loads(s)
 574             self.assertEqual(t, u)
 575             if hasattr(os, "stat"):
 576                 t = os.stat(os.curdir)
 577                 s = self.dumps(t, proto)
 578                 u = self.loads(s)
 579                 self.assertEqual(t, u)
 580             if hasattr(os, "statvfs"):
 581                 t = os.statvfs(os.curdir)
 582                 s = self.dumps(t, proto)
 583                 u = self.loads(s)
 584                 self.assertEqual(t, u)
 585
 586     # Tests for protocol 2
 587
 588     def test_proto(self):
 589         build_none = pickle.NONE + pickle.STOP
 590         for proto in protocols:
 591             expected = build_none
 592             if proto >= 2:
 593                 expected = pickle.PROTO + chr(proto) + expected
 594             p = self.dumps(None, proto)
 595             self.assertEqual(p, expected)
 596
 597         oob = protocols[-1] + 1     # a future protocol
 598         badpickle = pickle.PROTO + chr(oob) + build_none
 599         try:
 600             self.loads(badpickle)
 601         except ValueError, detail:
 602             self.failUnless(str(detail).startswith(
 603                                             "unsupported pickle protocol"))
 604         else:
 605             self.fail("expected bad protocol number to raise ValueError")
 606
 607     def test_long1(self):
 608         x = 12345678910111213141516178920L
 609         for proto in protocols:
 610             s = self.dumps(x, proto)
 611             y = self.loads(s)
 612             self.assertEqual(x, y)
 613             self.assertEqual(opcode_in_pickle(pickle.LONG1, s), proto >= 2)
 614
 615     def test_long4(self):
 616         x = 12345678910111213141516178920L << (256*8)
 617         for proto in protocols:
 618             s = self.dumps(x, proto)
 619             y = self.loads(s)
 620             self.assertEqual(x, y)
 621             self.assertEqual(opcode_in_pickle(pickle.LONG4, s), proto >= 2)
 622
 623     def test_short_tuples(self):
 624         # Map (proto, len(tuple)) to expected opcode.
 625         expected_opcode = {(0, 0): pickle.TUPLE,
 626                            (0, 1): pickle.TUPLE,
 627                            (0, 2): pickle.TUPLE,
 628                            (0, 3): pickle.TUPLE,
 629                            (0, 4): pickle.TUPLE,
 630
 631                            (1, 0): pickle.EMPTY_TUPLE,
 632                            (1, 1): pickle.TUPLE,
 633                            (1, 2): pickle.TUPLE,
 634                            (1, 3): pickle.TUPLE,
 635                            (1, 4): pickle.TUPLE,
 636
 637                            (2, 0): pickle.EMPTY_TUPLE,
 638                            (2, 1): pickle.TUPLE1,
 639                            (2, 2): pickle.TUPLE2,
 640                            (2, 3): pickle.TUPLE3,
 641                            (2, 4): pickle.TUPLE,
 642                           }
 643         a = ()
 644         b = (1,)
 645         c = (1, 2)
 646         d = (1, 2, 3)
 647         e = (1, 2, 3, 4)
 648         for proto in protocols:
 649             for x in a, b, c, d, e:
 650                 s = self.dumps(x, proto)
 651                 y = self.loads(s)
 652                 self.assertEqual(x, y, (proto, x, s, y))
 653                 expected = expected_opcode[proto, len(x)]
 654                 self.assertEqual(opcode_in_pickle(expected, s), True)
 655
 656     def test_singletons(self):
 657         # Map (proto, singleton) to expected opcode.
 658         expected_opcode = {(0, None): pickle.NONE,
 659                            (1, None): pickle.NONE,
 660                            (2, None): pickle.NONE,
 661
 662                            (0, True): pickle.INT,
 663                            (1, True): pickle.INT,
 664                            (2, True): pickle.NEWTRUE,
 665
 666                            (0, False): pickle.INT,
 667                            (1, False): pickle.INT,
 668                            (2, False): pickle.NEWFALSE,
 669                           }
 670         for proto in protocols:
 671             for x in None, False, True:
 672                 s = self.dumps(x, proto)
 673                 y = self.loads(s)
 674                 self.assert_(x is y, (proto, x, s, y))
 675                 expected = expected_opcode[proto, x]
 676                 self.assertEqual(opcode_in_pickle(expected, s), True)
 677
 678     def test_newobj_tuple(self):
 679         x = MyTuple([1, 2, 3])
 680         x.foo = 42
 681         x.bar = "hello"
 682         for proto in protocols:
 683             s = self.dumps(x, proto)
 684             y = self.loads(s)
 685             self.assertEqual(tuple(x), tuple(y))
 686             self.assertEqual(x.__dict__, y.__dict__)
 687
 688     def test_newobj_list(self):
 689         x = MyList([1, 2, 3])
 690         x.foo = 42
 691         x.bar = "hello"
 692         for proto in protocols:
 693             s = self.dumps(x, proto)
 694             y = self.loads(s)
 695             self.assertEqual(list(x), list(y))
 696             self.assertEqual(x.__dict__, y.__dict__)
 697
 698     def test_newobj_generic(self):
 699         for proto in protocols:
 700             for C in myclasses:
 701                 B = C.__base__
 702                 x = C(C.sample)
 703                 x.foo = 42
 704                 s = self.dumps(x, proto)
 705                 y = self.loads(s)
 706                 detail = (proto, C, B, x, y, type(y))
 707                 self.assertEqual(B(x), B(y), detail)
 708                 self.assertEqual(x.__dict__, y.__dict__, detail)
 709
 710     # Register a type with copy_reg, with extension code extcode.  Pickle
 711     # an object of that type.  Check that the resulting pickle uses opcode
 712     # (EXT[124]) under proto 2, and not in proto 1.
 713
 714     def produce_global_ext(self, extcode, opcode):
 715         e = ExtensionSaver(extcode)
 716         try:
 717             copy_reg.add_extension(__name__, "MyList", extcode)
 718             x = MyList([1, 2, 3])
 719             x.foo = 42
 720             x.bar = "hello"
 721
 722             # Dump using protocol 1 for comparison.
 723             s1 = self.dumps(x, 1)
 724             self.assert_(__name__ in s1)
 725             self.assert_("MyList" in s1)
 726             self.assertEqual(opcode_in_pickle(opcode, s1), False)
 727
 728             y = self.loads(s1)
 729             self.assertEqual(list(x), list(y))
 730             self.assertEqual(x.__dict__, y.__dict__)
 731
 732             # Dump using protocol 2 for test.
 733             s2 = self.dumps(x, 2)
 734             self.assert_(__name__ not in s2)
 735             self.assert_("MyList" not in s2)
 736             self.assertEqual(opcode_in_pickle(opcode, s2), True)
 737
 738             y = self.loads(s2)
 739             self.assertEqual(list(x), list(y))
 740             self.assertEqual(x.__dict__, y.__dict__)
 741
 742         finally:
 743             e.restore()
 744
 745     def test_global_ext1(self):
 746         self.produce_global_ext(0x00000001, pickle.EXT1)  # smallest EXT1 code
 747         self.produce_global_ext(0x000000ff, pickle.EXT1)  # largest EXT1 code
 748
 749     def test_global_ext2(self):
 750         self.produce_global_ext(0x00000100, pickle.EXT2)  # smallest EXT2 code
 751         self.produce_global_ext(0x0000ffff, pickle.EXT2)  # largest EXT2 code
 752         self.produce_global_ext(0x0000abcd, pickle.EXT2)  # check endianness
 753
 754     def test_global_ext4(self):
 755         self.produce_global_ext(0x00010000, pickle.EXT4)  # smallest EXT4 code
 756         self.produce_global_ext(0x7fffffff, pickle.EXT4)  # largest EXT4 code
 757         self.produce_global_ext(0x12abcdef, pickle.EXT4)  # check endianness
 758
 759     def test_list_chunking(self):
 760         n = 10  # too small to chunk
 761         x = range(n)
 762         for proto in protocols:
 763             s = self.dumps(x, proto)
 764             y = self.loads(s)
 765             self.assertEqual(x, y)
 766             num_appends = count_opcode(pickle.APPENDS, s)
 767             self.assertEqual(num_appends, proto > 0)
 768
 769         n = 2500  # expect at least two chunks when proto > 0
 770         x = range(n)
 771         for proto in protocols:
 772             s = self.dumps(x, proto)
 773             y = self.loads(s)
 774             self.assertEqual(x, y)
 775             num_appends = count_opcode(pickle.APPENDS, s)
 776             if proto == 0:
 777                 self.assertEqual(num_appends, 0)
 778             else:
 779                 self.failUnless(num_appends >= 2)
 780
 781     def test_dict_chunking(self):
 782         n = 10  # too small to chunk
 783         x = dict.fromkeys(range(n))
 784         for proto in protocols:
 785             s = self.dumps(x, proto)
 786             y = self.loads(s)
 787             self.assertEqual(x, y)
 788             num_setitems = count_opcode(pickle.SETITEMS, s)
 789             self.assertEqual(num_setitems, proto > 0)
 790
 791         n = 2500  # expect at least two chunks when proto > 0
 792         x = dict.fromkeys(range(n))
 793         for proto in protocols:
 794             s = self.dumps(x, proto)
 795             y = self.loads(s)
 796             self.assertEqual(x, y)
 797             num_setitems = count_opcode(pickle.SETITEMS, s)
 798             if proto == 0:
 799                 self.assertEqual(num_setitems, 0)
 800             else:
 801                 self.failUnless(num_setitems >= 2)
 802
 803     def test_simple_newobj(self):
 804         x = object.__new__(SimpleNewObj)  # avoid __init__
 805         x.abc = 666
 806         for proto in protocols:
 807             s = self.dumps(x, proto)
 808             self.assertEqual(opcode_in_pickle(pickle.NEWOBJ, s), proto >= 2)
 809             y = self.loads(s)   # will raise TypeError if __init__ called
 810             self.assertEqual(y.abc, 666)
 811             self.assertEqual(x.__dict__, y.__dict__)
 812
 813     def test_newobj_list_slots(self):
 814         x = SlotList([1, 2, 3])
 815         x.foo = 42
 816         x.bar = "hello"
 817         s = self.dumps(x, 2)
 818         y = self.loads(s)
 819         self.assertEqual(list(x), list(y))
 820         self.assertEqual(x.__dict__, y.__dict__)
 821         self.assertEqual(x.foo, y.foo)
 822         self.assertEqual(x.bar, y.bar)
 823
 824     def test_reduce_overrides_default_reduce_ex(self):
 825         for proto in 0, 1, 2:
 826             x = REX_one()
 827             self.assertEqual(x._reduce_called, 0)
 828             s = self.dumps(x, proto)
 829             self.assertEqual(x._reduce_called, 1)
 830             y = self.loads(s)
 831             self.assertEqual(y._reduce_called, 0)
 832
 833     def test_reduce_ex_called(self):
 834         for proto in 0, 1, 2:
 835             x = REX_two()
 836             self.assertEqual(x._proto, None)
 837             s = self.dumps(x, proto)
 838             self.assertEqual(x._proto, proto)
 839             y = self.loads(s)
 840             self.assertEqual(y._proto, None)
 841
 842     def test_reduce_ex_overrides_reduce(self):
 843         for proto in 0, 1, 2:
 844             x = REX_three()
 845             self.assertEqual(x._proto, None)
 846             s = self.dumps(x, proto)
 847             self.assertEqual(x._proto, proto)
 848             y = self.loads(s)
 849             self.assertEqual(y._proto, None)
 850
 851     def test_reduce_ex_calls_base(self):
 852         for proto in 0, 1, 2:
 853             x = REX_four()
 854             self.assertEqual(x._proto, None)
 855             s = self.dumps(x, proto)
 856             self.assertEqual(x._proto, proto)
 857             y = self.loads(s)
 858             self.assertEqual(y._proto, proto)
 859
 860     def test_reduce_calls_base(self):
 861         for proto in 0, 1, 2:
 862             x = REX_five()
 863             self.assertEqual(x._reduce_called, 0)
 864             s = self.dumps(x, proto)
 865             self.assertEqual(x._reduce_called, 1)
 866             y = self.loads(s)
 867             self.assertEqual(y._reduce_called, 1)
 868
 869     def test_reduce_bad_iterator(self):
 870         # Issue4176: crash when 4th and 5th items of __reduce__()
 871         # are not iterators
 872         class C(object):
 873             def __reduce__(self):
 874                 # 4th item is not an iterator
 875                 return list, (), None, [], None
 876         class D(object):
 877             def __reduce__(self):
 878                 # 5th item is not an iterator
 879                 return dict, (), None, None, []
 880
 881         # Protocol 0 is less strict and also accept iterables.
 882         for proto in 0, 1, 2:
 883             try:
 884                 self.dumps(C(), proto)
 885             except (AttributeError, pickle.PickleError, cPickle.PickleError):
 886                 pass
 887             try:
 888                 self.dumps(D(), proto)
 889             except (AttributeError, pickle.PickleError, cPickle.PickleError):
 890                 pass
 891
 892 # Test classes for reduce_ex
 893
 894 class REX_one(object):
 895     _reduce_called = 0
 896     def __reduce__(self):
 897         self._reduce_called = 1
 898         return REX_one, ()
 899     # No __reduce_ex__ here, but inheriting it from object
 900
 901 class REX_two(object):
 902     _proto = None
 903     def __reduce_ex__(self, proto):
 904         self._proto = proto
 905         return REX_two, ()
 906     # No __reduce__ here, but inheriting it from object
 907
 908 class REX_three(object):
 909     _proto = None
 910     def __reduce_ex__(self, proto):
 911         self._proto = proto
 912         return REX_two, ()
 913     def __reduce__(self):
 914         raise TestFailed, "This __reduce__ shouldn't be called"
 915
 916 class REX_four(object):
 917     _proto = None
 918     def __reduce_ex__(self, proto):
 919         self._proto = proto
 920         return object.__reduce_ex__(self, proto)
 921     # Calling base class method should succeed
 922
 923 class REX_five(object):
 924     _reduce_called = 0
 925     def __reduce__(self):
 926         self._reduce_called = 1
 927         return object.__reduce__(self)
 928     # This one used to fail with infinite recursion
 929
 930 # Test classes for newobj
 931
 932 class MyInt(int):
 933     sample = 1
 934
 935 class MyLong(long):
 936     sample = 1L
 937
 938 class MyFloat(float):
 939     sample = 1.0
 940
 941 class MyComplex(complex):
 942     sample = 1.0 + 0.0j
 943
 944 class MyStr(str):
 945     sample = "hello"
 946
 947 class MyUnicode(unicode):
 948     sample = u"hello \u1234"
 949
 950 class MyTuple(tuple):
 951     sample = (1, 2, 3)
 952
 953 class MyList(list):
 954     sample = [1, 2, 3]
 955
 956 class MyDict(dict):
 957     sample = {"a": 1, "b": 2}
 958
 959 myclasses = [MyInt, MyLong, MyFloat,
 960              MyComplex,
 961              MyStr, MyUnicode,
 962              MyTuple, MyList, MyDict]
 963
 964
 965 class SlotList(MyList):
 966     __slots__ = ["foo"]
 967
 968 class SimpleNewObj(object):
 969     def __init__(self, a, b, c):
 970         # raise an error, to make sure this isn't called
 971         raise TypeError("SimpleNewObj.__init__() didn't expect to get called")
 972
 973 class AbstractPickleModuleTests(unittest.TestCase):
 974
 975     def test_dump_closed_file(self):
 976         import os
 977         f = open(TESTFN, "w")
 978         try:
 979             f.close()
 980             self.assertRaises(ValueError, self.module.dump, 123, f)
 981         finally:
 982             os.remove(TESTFN)
 983
 984     def test_load_closed_file(self):
 985         import os
 986         f = open(TESTFN, "w")
 987         try:
 988             f.close()
 989             self.assertRaises(ValueError, self.module.dump, 123, f)
 990         finally:
 991             os.remove(TESTFN)
 992
 993     def test_highest_protocol(self):
 994         # Of course this needs to be changed when HIGHEST_PROTOCOL changes.
 995         self.assertEqual(self.module.HIGHEST_PROTOCOL, 2)
 996
 997     def test_callapi(self):
 998         from cStringIO import StringIO
 999         f = StringIO()
1000         # With and without keyword arguments
1001         self.module.dump(123, f, -1)
1002         self.module.dump(123, file=f, protocol=-1)
1003         self.module.dumps(123, -1)
1004         self.module.dumps(123, protocol=-1)
1005         self.module.Pickler(f, -1)
1006         self.module.Pickler(f, protocol=-1)
1007
1008 class AbstractPersistentPicklerTests(unittest.TestCase):
1009
1010     # This class defines persistent_id() and persistent_load()
1011     # functions that should be used by the pickler.  All even integers
1012     # are pickled using persistent ids.
1013
1014     def persistent_id(self, object):
1015         if isinstance(object, int) and object % 2 == 0:
1016             self.id_count += 1
1017             return str(object)
1018         else:
1019             return None
1020
1021     def persistent_load(self, oid):
1022         self.load_count += 1
1023         object = int(oid)
1024         assert object % 2 == 0
1025         return object
1026
1027     def test_persistence(self):
1028         self.id_count = 0
1029         self.load_count = 0
1030         L = range(10)
1031         self.assertEqual(self.loads(self.dumps(L)), L)
1032         self.assertEqual(self.id_count, 5)
1033         self.assertEqual(self.load_count, 5)
1034
1035     def test_bin_persistence(self):
1036         self.id_count = 0
1037         self.load_count = 0
1038         L = range(10)
1039         self.assertEqual(self.loads(self.dumps(L, 1)), L)
1040         self.assertEqual(self.id_count, 5)
1041         self.assertEqual(self.load_count, 5)