test_pickle.py revision 0c5958b1636c47ed7c284f859c8e805fd06a0e63 
1import pickle
2from cStringIO import StringIO
3
4from test import test_support
5
6from test.pickletester import (AbstractPickleTests,
7                               AbstractPickleModuleTests,
8                               AbstractPersistentPicklerTests,
9                               AbstractPicklerUnpicklerObjectTests,
10                               BigmemPickleTests)
11
12class PickleTests(AbstractPickleTests, AbstractPickleModuleTests):
13
14    def dumps(self, arg, proto=0, fast=0):
15        # Ignore fast
16        return pickle.dumps(arg, proto)
17
18    def loads(self, buf):
19        # Ignore fast
20        return pickle.loads(buf)
21
22    module = pickle
23    error = KeyError
24
25class PicklerTests(AbstractPickleTests):
26
27    error = KeyError
28
29    def dumps(self, arg, proto=0, fast=0):
30        f = StringIO()
31        p = pickle.Pickler(f, proto)
32        if fast:
33            p.fast = fast
34        p.dump(arg)
35        f.seek(0)
36        return f.read()
37
38    def loads(self, buf):
39        f = StringIO(buf)
40        u = pickle.Unpickler(f)
41        return u.load()
42
43class PersPicklerTests(AbstractPersistentPicklerTests):
44
45    def dumps(self, arg, proto=0, fast=0):
46        class PersPickler(pickle.Pickler):
47            def persistent_id(subself, obj):
48                return self.persistent_id(obj)
49        f = StringIO()
50        p = PersPickler(f, proto)
51        if fast:
52            p.fast = fast
53        p.dump(arg)
54        f.seek(0)
55        return f.read()
56
57    def loads(self, buf):
58        class PersUnpickler(pickle.Unpickler):
59            def persistent_load(subself, obj):
60                return self.persistent_load(obj)
61        f = StringIO(buf)
62        u = PersUnpickler(f)
63        return u.load()
64
65class PicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
66
67    pickler_class = pickle.Pickler
68    unpickler_class = pickle.Unpickler
69
70class PickleBigmemPickleTests(BigmemPickleTests):
71
72    def dumps(self, arg, proto=0, fast=0):
73        # Ignore fast
74        return pickle.dumps(arg, proto)
75
76    def loads(self, buf):
77        # Ignore fast
78        return pickle.loads(buf)
79
80
81def test_main():
82    test_support.run_unittest(
83        PickleTests,
84        PicklerTests,
85        PersPicklerTests,
86        PicklerUnpicklerObjectTests,
87        PickleBigmemPickleTests,
88    )
89    test_support.run_doctest(pickle)
90
91if __name__ == "__main__":
92    test_main()
93