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