1 import cPickle, unittest

     2 from cStringIO import StringIO

     3 from test.pickletester import AbstractPickleTests, AbstractPickleModuleTests

     4 from test import test_support

     5 

     6 class cPickleTests(AbstractPickleTests, AbstractPickleModuleTests):

     7 

     8     def setUp(self):

     9         self.dumps = cPickle.dumps

    10         self.loads = cPickle.loads

    11 

    12     error = cPickle.BadPickleGet

    13     module = cPickle

    14 

    15 class cPicklePicklerTests(AbstractPickleTests):

    16 

    17     def dumps(self, arg, proto=0):

    18         f = StringIO()

    19         p = cPickle.Pickler(f, proto)

    20         p.dump(arg)

    21         f.seek(0)

    22         return f.read()

    23 

    24     def loads(self, buf):

    25         f = StringIO(buf)

    26         p = cPickle.Unpickler(f)

    27         return p.load()

    28 

    29     error = cPickle.BadPickleGet

    30 

    31 class cPickleListPicklerTests(AbstractPickleTests):

    32 

    33     def dumps(self, arg, proto=0):

    34         p = cPickle.Pickler(proto)

    35         p.dump(arg)

    36         return p.getvalue()

    37 

    38     def loads(self, *args):

    39         f = StringIO(args[0])

    40         p = cPickle.Unpickler(f)

    41         return p.load()

    42 

    43     error = cPickle.BadPickleGet

    44 

    45 class cPickleFastPicklerTests(AbstractPickleTests):

    46 

    47     def dumps(self, arg, proto=0):

    48         f = StringIO()

    49         p = cPickle.Pickler(f, proto)

    50         p.fast = 1

    51         p.dump(arg)

    52         f.seek(0)

    53         return f.read()

    54 

    55     def loads(self, *args):

    56         f = StringIO(args[0])

    57         p = cPickle.Unpickler(f)

    58         return p.load()

    59 

    60     error = cPickle.BadPickleGet

    61 

    62     def test_recursive_list(self):

    63         self.assertRaises(ValueError,

    64                           AbstractPickleTests.test_recursive_list,

    65                           self)

    66 

    67     def test_recursive_inst(self):

    68         self.assertRaises(ValueError,

    69                           AbstractPickleTests.test_recursive_inst,

    70                           self)

    71 

    72     def test_recursive_dict(self):

    73         self.assertRaises(ValueError,

    74                           AbstractPickleTests.test_recursive_dict,

    75                           self)

    76 

    77     def test_recursive_multi(self):

    78         self.assertRaises(ValueError,

    79                           AbstractPickleTests.test_recursive_multi,

    80                           self)

    81 

    82     def test_nonrecursive_deep(self):

    83         # If it's not cyclic, it should pickle OK even if the nesting

    84         # depth exceeds PY_CPICKLE_FAST_LIMIT.  That happens to be

    85         # 50 today.  Jack Jansen reported stack overflow on Mac OS 9

    86         # at 64.

    87         a = []

    88         for i in range(60):

    89             a = [a]

    90         b = self.loads(self.dumps(a))

    91         self.assertEqual(a, b)

    92 

    93 class Node(object):

    94     pass

    95 

    96 class cPickleDeepRecursive(unittest.TestCase):

    97     def test_issue2702(self):

    98         # This should raise a RecursionLimit but in some

    99         # platforms (FreeBSD, win32) sometimes raises KeyError instead,

   100         # or just silently terminates the interpreter (=crashes).

   101         nodes = [Node() for i in range(500)]

   102         for n in nodes:

   103             n.connections = list(nodes)

   104             n.connections.remove(n)

   105         self.assertRaises(RuntimeError, cPickle.dumps, n)

   106 

   107     def test_issue3179(self):

   108         # Safe test, because I broke this case when fixing the

   109         # behaviour for the previous test.

   110         res=[]

   111         for x in range(1,2000):

   112             res.append(dict(doc=x, similar=[]))

   113         cPickle.dumps(res)

   114 

   115 

   116 def test_main():

   117     test_support.run_unittest(

   118         cPickleTests,

   119         cPicklePicklerTests,

   120         cPickleListPicklerTests,

   121         cPickleFastPicklerTests,

   122         cPickleDeepRecursive,

   123     )

   124 

   125 if __name__ == "__main__":

   126     test_main()