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