Jeremy Hylton | be467e5 | 2000-09-15 15:14:51 +0000 | [diff] [blame] | 1 | import pickle |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 2 | import io |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 3 | import collections |
Tim Peters | 47a6b13 | 2003-01-28 22:34:11 +0000 | [diff] [blame] | 4 | |
Benjamin Peterson | ee8712c | 2008-05-20 21:35:26 +0000 | [diff] [blame] | 5 | from test import support |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 6 | |
Tim Peters | 47a6b13 | 2003-01-28 22:34:11 +0000 | [diff] [blame] | 7 | from test.pickletester import AbstractPickleTests |
Tim Peters | 47a6b13 | 2003-01-28 22:34:11 +0000 | [diff] [blame] | 8 | from test.pickletester import AbstractPickleModuleTests |
| 9 | from test.pickletester import AbstractPersistentPicklerTests |
Collin Winter | 771d834 | 2009-04-16 03:18:06 +0000 | [diff] [blame] | 10 | from test.pickletester import AbstractPicklerUnpicklerObjectTests |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 11 | from test.pickletester import AbstractDispatchTableTests |
Antoine Pitrou | 82be19f | 2011-08-29 23:09:33 +0200 | [diff] [blame] | 12 | from test.pickletester import BigmemPickleTests |
Tim Peters | 47a6b13 | 2003-01-28 22:34:11 +0000 | [diff] [blame] | 13 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 14 | try: |
| 15 | import _pickle |
| 16 | has_c_implementation = True |
| 17 | except ImportError: |
| 18 | has_c_implementation = False |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 19 | |
Guido van Rossum | 98297ee | 2007-11-06 21:34:58 +0000 | [diff] [blame] | 20 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 21 | class PickleTests(AbstractPickleModuleTests): |
| 22 | pass |
Guido van Rossum | 5d9113d | 2003-01-29 17:58:45 +0000 | [diff] [blame] | 23 | |
Tim Peters | e0c446b | 2001-10-18 21:57:37 +0000 | [diff] [blame] | 24 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 25 | class PyPicklerTests(AbstractPickleTests): |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 26 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 27 | pickler = pickle._Pickler |
| 28 | unpickler = pickle._Unpickler |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 29 | |
Guido van Rossum | f416981 | 2008-03-17 22:56:06 +0000 | [diff] [blame] | 30 | def dumps(self, arg, proto=None): |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 31 | f = io.BytesIO() |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 32 | p = self.pickler(f, proto) |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 33 | p.dump(arg) |
| 34 | f.seek(0) |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 35 | return bytes(f.read()) |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 36 | |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 37 | def loads(self, buf, **kwds): |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 38 | f = io.BytesIO(buf) |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 39 | u = self.unpickler(f, **kwds) |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 40 | return u.load() |
| 41 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 42 | |
Antoine Pitrou | 82be19f | 2011-08-29 23:09:33 +0200 | [diff] [blame] | 43 | class InMemoryPickleTests(AbstractPickleTests, BigmemPickleTests): |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 44 | |
| 45 | pickler = pickle._Pickler |
| 46 | unpickler = pickle._Unpickler |
| 47 | |
Antoine Pitrou | 82be19f | 2011-08-29 23:09:33 +0200 | [diff] [blame] | 48 | def dumps(self, arg, protocol=None): |
| 49 | return pickle.dumps(arg, protocol) |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 50 | |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 51 | def loads(self, buf, **kwds): |
| 52 | return pickle.loads(buf, **kwds) |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 53 | |
| 54 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 55 | class PyPersPicklerTests(AbstractPersistentPicklerTests): |
| 56 | |
| 57 | pickler = pickle._Pickler |
| 58 | unpickler = pickle._Unpickler |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 59 | |
Guido van Rossum | f416981 | 2008-03-17 22:56:06 +0000 | [diff] [blame] | 60 | def dumps(self, arg, proto=None): |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 61 | class PersPickler(self.pickler): |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 62 | def persistent_id(subself, obj): |
| 63 | return self.persistent_id(obj) |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 64 | f = io.BytesIO() |
Guido van Rossum | 9d32bb1 | 2003-01-28 03:51:53 +0000 | [diff] [blame] | 65 | p = PersPickler(f, proto) |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 66 | p.dump(arg) |
| 67 | f.seek(0) |
| 68 | return f.read() |
| 69 | |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 70 | def loads(self, buf, **kwds): |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 71 | class PersUnpickler(self.unpickler): |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 72 | def persistent_load(subself, obj): |
| 73 | return self.persistent_load(obj) |
Guido van Rossum | cfe5f20 | 2007-05-08 21:26:54 +0000 | [diff] [blame] | 74 | f = io.BytesIO(buf) |
Alexander Belopolsky | ec8f0df | 2011-02-24 20:34:38 +0000 | [diff] [blame] | 75 | u = PersUnpickler(f, **kwds) |
Jeremy Hylton | 5e0f4e7 | 2002-11-13 22:01:27 +0000 | [diff] [blame] | 76 | return u.load() |
| 77 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 78 | |
Collin Winter | 771d834 | 2009-04-16 03:18:06 +0000 | [diff] [blame] | 79 | class PyPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests): |
| 80 | |
| 81 | pickler_class = pickle._Pickler |
| 82 | unpickler_class = pickle._Unpickler |
| 83 | |
| 84 | |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 85 | class PyDispatchTableTests(AbstractDispatchTableTests): |
| 86 | pickler_class = pickle._Pickler |
| 87 | def get_dispatch_table(self): |
| 88 | return pickle.dispatch_table.copy() |
| 89 | |
| 90 | |
| 91 | class PyChainDispatchTableTests(AbstractDispatchTableTests): |
| 92 | pickler_class = pickle._Pickler |
| 93 | def get_dispatch_table(self): |
| 94 | return collections.ChainMap({}, pickle.dispatch_table) |
| 95 | |
| 96 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 97 | if has_c_implementation: |
| 98 | class CPicklerTests(PyPicklerTests): |
| 99 | pickler = _pickle.Pickler |
| 100 | unpickler = _pickle.Unpickler |
| 101 | |
| 102 | class CPersPicklerTests(PyPersPicklerTests): |
| 103 | pickler = _pickle.Pickler |
| 104 | unpickler = _pickle.Unpickler |
| 105 | |
Collin Winter | 771d834 | 2009-04-16 03:18:06 +0000 | [diff] [blame] | 106 | class CDumpPickle_LoadPickle(PyPicklerTests): |
| 107 | pickler = _pickle.Pickler |
| 108 | unpickler = pickle._Unpickler |
| 109 | |
| 110 | class DumpPickle_CLoadPickle(PyPicklerTests): |
| 111 | pickler = pickle._Pickler |
| 112 | unpickler = _pickle.Unpickler |
| 113 | |
| 114 | class CPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests): |
| 115 | pickler_class = _pickle.Pickler |
| 116 | unpickler_class = _pickle.Unpickler |
| 117 | |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 118 | class CDispatchTableTests(AbstractDispatchTableTests): |
| 119 | pickler_class = pickle.Pickler |
| 120 | def get_dispatch_table(self): |
| 121 | return pickle.dispatch_table.copy() |
| 122 | |
| 123 | class CChainDispatchTableTests(AbstractDispatchTableTests): |
| 124 | pickler_class = pickle.Pickler |
| 125 | def get_dispatch_table(self): |
| 126 | return collections.ChainMap({}, pickle.dispatch_table) |
| 127 | |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 128 | |
Fred Drake | 694ed09 | 2001-12-19 16:42:15 +0000 | [diff] [blame] | 129 | def test_main(): |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 130 | tests = [PickleTests, PyPicklerTests, PyPersPicklerTests, |
| 131 | PyDispatchTableTests, PyChainDispatchTableTests] |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 132 | if has_c_implementation: |
Collin Winter | 771d834 | 2009-04-16 03:18:06 +0000 | [diff] [blame] | 133 | tests.extend([CPicklerTests, CPersPicklerTests, |
| 134 | CDumpPickle_LoadPickle, DumpPickle_CLoadPickle, |
| 135 | PyPicklerUnpicklerObjectTests, |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 136 | CPicklerUnpicklerObjectTests, |
Antoine Pitrou | 8d3c290 | 2012-03-04 18:31:48 +0100 | [diff] [blame] | 137 | CDispatchTableTests, CChainDispatchTableTests, |
Antoine Pitrou | ea99c5c | 2010-09-09 18:33:21 +0000 | [diff] [blame] | 138 | InMemoryPickleTests]) |
Alexandre Vassalotti | ca2d610 | 2008-06-12 18:26:05 +0000 | [diff] [blame] | 139 | support.run_unittest(*tests) |
Benjamin Peterson | ee8712c | 2008-05-20 21:35:26 +0000 | [diff] [blame] | 140 | support.run_doctest(pickle) |
Fred Drake | 694ed09 | 2001-12-19 16:42:15 +0000 | [diff] [blame] | 141 | |
Jeremy Hylton | 6642653 | 2001-10-15 21:38:56 +0000 | [diff] [blame] | 142 | if __name__ == "__main__": |
Fred Drake | 694ed09 | 2001-12-19 16:42:15 +0000 | [diff] [blame] | 143 | test_main() |