[Python-checkins] cpython: Issue #14166: Pickler objects now have an optional `dispatch_table` attribute
antoine.pitrou
python-checkins at python.org
Sun Mar 4 18:35:49 CET 2012
http://hg.python.org/cpython/rev/f7a9a10ae0c0
changeset: 75391:f7a9a10ae0c0
user: Antoine Pitrou <solipsis at pitrou.net>
date: Sun Mar 04 18:31:48 2012 +0100
summary:
Issue #14166: Pickler objects now have an optional `dispatch_table` attribute which allows to set custom per-pickler reduction functions.
Patch by sbt.
files:
Doc/library/copyreg.rst | 8 +-
Doc/library/pickle.rst | 61 +++++++++++++++++
Lib/pickle.py | 4 +-
Lib/test/pickletester.py | 99 ++++++++++++++++++++++++++++
Lib/test/test_pickle.py | 28 +++++++-
Misc/NEWS | 4 +
Modules/_pickle.c | 39 ++++++++--
7 files changed, 228 insertions(+), 15 deletions(-)
diff --git a/Doc/library/copyreg.rst b/Doc/library/copyreg.rst
--- a/Doc/library/copyreg.rst
+++ b/Doc/library/copyreg.rst
@@ -32,6 +32,8 @@
returned by *function* at pickling time. :exc:`TypeError` will be raised if
*object* is a class or *constructor* is not callable.
- See the :mod:`pickle` module for more details on the interface expected of
- *function* and *constructor*.
-
+ See the :mod:`pickle` module for more details on the interface
+ expected of *function* and *constructor*. Note that the
+ :attr:`~pickle.Pickler.dispatch_table` attribute of a pickler
+ object or subclass of :class:`pickle.Pickler` can also be used for
+ declaring reduction functions.
diff --git a/Doc/library/pickle.rst b/Doc/library/pickle.rst
--- a/Doc/library/pickle.rst
+++ b/Doc/library/pickle.rst
@@ -285,6 +285,29 @@
See :ref:`pickle-persistent` for details and examples of uses.
+ .. attribute:: dispatch_table
+
+ A pickler object's dispatch table is a registry of *reduction
+ functions* of the kind which can be declared using
+ :func:`copyreg.pickle`. It is a mapping whose keys are classes
+ and whose values are reduction functions. A reduction function
+ takes a single argument of the associated class and should
+ conform to the same interface as a :meth:`~object.__reduce__`
+ method.
+
+ By default, a pickler object will not have a
+ :attr:`dispatch_table` attribute, and it will instead use the
+ global dispatch table managed by the :mod:`copyreg` module.
+ However, to customize the pickling for a specific pickler object
+ one can set the :attr:`dispatch_table` attribute to a dict-like
+ object. Alternatively, if a subclass of :class:`Pickler` has a
+ :attr:`dispatch_table` attribute then this will be used as the
+ default dispatch table for instances of that class.
+
+ See :ref:`pickle-dispatch` for usage examples.
+
+ .. versionadded:: 3.3
+
.. attribute:: fast
Deprecated. Enable fast mode if set to a true value. The fast mode
@@ -575,6 +598,44 @@
.. literalinclude:: ../includes/dbpickle.py
+.. _pickle-dispatch:
+
+Dispatch Tables
+^^^^^^^^^^^^^^^
+
+If one wants to customize pickling of some classes without disturbing
+any other code which depends on pickling, then one can create a
+pickler with a private dispatch table.
+
+The global dispatch table managed by the :mod:`copyreg` module is
+available as :data:`copyreg.dispatch_table`. Therefore, one may
+choose to use a modified copy of :data:`copyreg.dispatch_table` as a
+private dispatch table.
+
+For example ::
+
+ f = io.BytesIO()
+ p = pickle.Pickler(f)
+ p.dispatch_table = copyreg.dispatch_table.copy()
+ p.dispatch_table[SomeClass] = reduce_SomeClass
+
+creates an instance of :class:`pickle.Pickler` with a private dispatch
+table which handles the ``SomeClass`` class specially. Alternatively,
+the code ::
+
+ class MyPickler(pickle.Pickler):
+ dispatch_table = copyreg.dispatch_table.copy()
+ dispatch_table[SomeClass] = reduce_SomeClass
+ f = io.BytesIO()
+ p = MyPickler(f)
+
+does the same, but all instances of ``MyPickler`` will by default
+share the same dispatch table. The equivalent code using the
+:mod:`copyreg` module is ::
+
+ copyreg.pickle(SomeClass, reduce_SomeClass)
+ f = io.BytesIO()
+ p = pickle.Pickler(f)
.. _pickle-state:
diff --git a/Lib/pickle.py b/Lib/pickle.py
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -297,8 +297,8 @@
f(self, obj) # Call unbound method with explicit self
return
- # Check copyreg.dispatch_table
- reduce = dispatch_table.get(t)
+ # Check private dispatch table if any, or else copyreg.dispatch_table
+ reduce = getattr(self, 'dispatch_table', dispatch_table).get(t)
if reduce:
rv = reduce(obj)
else:
diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -1605,6 +1605,105 @@
self.assertEqual(unpickler.load(), data)
+# Tests for dispatch_table attribute
+
+REDUCE_A = 'reduce_A'
+
+class AAA(object):
+ def __reduce__(self):
+ return str, (REDUCE_A,)
+
+class BBB(object):
+ pass
+
+class AbstractDispatchTableTests(unittest.TestCase):
+
+ def test_default_dispatch_table(self):
+ # No dispatch_table attribute by default
+ f = io.BytesIO()
+ p = self.pickler_class(f, 0)
+ with self.assertRaises(AttributeError):
+ p.dispatch_table
+ self.assertFalse(hasattr(p, 'dispatch_table'))
+
+ def test_class_dispatch_table(self):
+ # A dispatch_table attribute can be specified class-wide
+ dt = self.get_dispatch_table()
+
+ class MyPickler(self.pickler_class):
+ dispatch_table = dt
+
+ def dumps(obj, protocol=None):
+ f = io.BytesIO()
+ p = MyPickler(f, protocol)
+ self.assertEqual(p.dispatch_table, dt)
+ p.dump(obj)
+ return f.getvalue()
+
+ self._test_dispatch_table(dumps, dt)
+
+ def test_instance_dispatch_table(self):
+ # A dispatch_table attribute can also be specified instance-wide
+ dt = self.get_dispatch_table()
+
+ def dumps(obj, protocol=None):
+ f = io.BytesIO()
+ p = self.pickler_class(f, protocol)
+ p.dispatch_table = dt
+ self.assertEqual(p.dispatch_table, dt)
+ p.dump(obj)
+ return f.getvalue()
+
+ self._test_dispatch_table(dumps, dt)
+
+ def _test_dispatch_table(self, dumps, dispatch_table):
+ def custom_load_dump(obj):
+ return pickle.loads(dumps(obj, 0))
+
+ def default_load_dump(obj):
+ return pickle.loads(pickle.dumps(obj, 0))
+
+ # pickling complex numbers using protocol 0 relies on copyreg
+ # so check pickling a complex number still works
+ z = 1 + 2j
+ self.assertEqual(custom_load_dump(z), z)
+ self.assertEqual(default_load_dump(z), z)
+
+ # modify pickling of complex
+ REDUCE_1 = 'reduce_1'
+ def reduce_1(obj):
+ return str, (REDUCE_1,)
+ dispatch_table[complex] = reduce_1
+ self.assertEqual(custom_load_dump(z), REDUCE_1)
+ self.assertEqual(default_load_dump(z), z)
+
+ # check picklability of AAA and BBB
+ a = AAA()
+ b = BBB()
+ self.assertEqual(custom_load_dump(a), REDUCE_A)
+ self.assertIsInstance(custom_load_dump(b), BBB)
+ self.assertEqual(default_load_dump(a), REDUCE_A)
+ self.assertIsInstance(default_load_dump(b), BBB)
+
+ # modify pickling of BBB
+ dispatch_table[BBB] = reduce_1
+ self.assertEqual(custom_load_dump(a), REDUCE_A)
+ self.assertEqual(custom_load_dump(b), REDUCE_1)
+ self.assertEqual(default_load_dump(a), REDUCE_A)
+ self.assertIsInstance(default_load_dump(b), BBB)
+
+ # revert pickling of BBB and modify pickling of AAA
+ REDUCE_2 = 'reduce_2'
+ def reduce_2(obj):
+ return str, (REDUCE_2,)
+ dispatch_table[AAA] = reduce_2
+ del dispatch_table[BBB]
+ self.assertEqual(custom_load_dump(a), REDUCE_2)
+ self.assertIsInstance(custom_load_dump(b), BBB)
+ self.assertEqual(default_load_dump(a), REDUCE_A)
+ self.assertIsInstance(default_load_dump(b), BBB)
+
+
if __name__ == "__main__":
# Print some stuff that can be used to rewrite DATA{0,1,2}
from pickletools import dis
diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py
--- a/Lib/test/test_pickle.py
+++ b/Lib/test/test_pickle.py
@@ -1,5 +1,6 @@
import pickle
import io
+import collections
from test import support
@@ -7,6 +8,7 @@
from test.pickletester import AbstractPickleModuleTests
from test.pickletester import AbstractPersistentPicklerTests
from test.pickletester import AbstractPicklerUnpicklerObjectTests
+from test.pickletester import AbstractDispatchTableTests
from test.pickletester import BigmemPickleTests
try:
@@ -80,6 +82,18 @@
unpickler_class = pickle._Unpickler
+class PyDispatchTableTests(AbstractDispatchTableTests):
+ pickler_class = pickle._Pickler
+ def get_dispatch_table(self):
+ return pickle.dispatch_table.copy()
+
+
+class PyChainDispatchTableTests(AbstractDispatchTableTests):
+ pickler_class = pickle._Pickler
+ def get_dispatch_table(self):
+ return collections.ChainMap({}, pickle.dispatch_table)
+
+
if has_c_implementation:
class CPicklerTests(PyPicklerTests):
pickler = _pickle.Pickler
@@ -101,14 +115,26 @@
pickler_class = _pickle.Pickler
unpickler_class = _pickle.Unpickler
+ class CDispatchTableTests(AbstractDispatchTableTests):
+ pickler_class = pickle.Pickler
+ def get_dispatch_table(self):
+ return pickle.dispatch_table.copy()
+
+ class CChainDispatchTableTests(AbstractDispatchTableTests):
+ pickler_class = pickle.Pickler
+ def get_dispatch_table(self):
+ return collections.ChainMap({}, pickle.dispatch_table)
+
def test_main():
- tests = [PickleTests, PyPicklerTests, PyPersPicklerTests]
+ tests = [PickleTests, PyPicklerTests, PyPersPicklerTests,
+ PyDispatchTableTests, PyChainDispatchTableTests]
if has_c_implementation:
tests.extend([CPicklerTests, CPersPicklerTests,
CDumpPickle_LoadPickle, DumpPickle_CLoadPickle,
PyPicklerUnpicklerObjectTests,
CPicklerUnpicklerObjectTests,
+ CDispatchTableTests, CChainDispatchTableTests,
InMemoryPickleTests])
support.run_unittest(*tests)
support.run_doctest(pickle)
diff --git a/Misc/NEWS b/Misc/NEWS
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -511,6 +511,10 @@
Library
-------
+- Issue #14166: Pickler objects now have an optional ``dispatch_table``
+ attribute which allows to set custom per-pickler reduction functions.
+ Patch by sbt.
+
- Issue #14177: marshal.loads() now raises TypeError when given an unicode
string. Patch by Guilherme Gonçalves.
diff --git a/Modules/_pickle.c b/Modules/_pickle.c
--- a/Modules/_pickle.c
+++ b/Modules/_pickle.c
@@ -319,6 +319,7 @@
objects to support self-referential objects
pickling. */
PyObject *pers_func; /* persistent_id() method, can be NULL */
+ PyObject *dispatch_table; /* private dispatch_table, can be NULL */
PyObject *arg;
PyObject *write; /* write() method of the output stream. */
@@ -764,6 +765,7 @@
return NULL;
self->pers_func = NULL;
+ self->dispatch_table = NULL;
self->arg = NULL;
self->write = NULL;
self->proto = 0;
@@ -3176,17 +3178,24 @@
/* XXX: This part needs some unit tests. */
/* Get a reduction callable, and call it. This may come from
- * copyreg.dispatch_table, the object's __reduce_ex__ method,
- * or the object's __reduce__ method.
+ * self.dispatch_table, copyreg.dispatch_table, the object's
+ * __reduce_ex__ method, or the object's __reduce__ method.
*/
- reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
+ if (self->dispatch_table == NULL) {
+ reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
+ /* PyDict_GetItem() unlike PyObject_GetItem() and
+ PyObject_GetAttr() returns a borrowed ref */
+ Py_XINCREF(reduce_func);
+ } else {
+ reduce_func = PyObject_GetItem(self->dispatch_table, (PyObject *)type);
+ if (reduce_func == NULL) {
+ if (PyErr_ExceptionMatches(PyExc_KeyError))
+ PyErr_Clear();
+ else
+ goto error;
+ }
+ }
if (reduce_func != NULL) {
- /* Here, the reference count of the reduce_func object returned by
- PyDict_GetItem needs to be increased to be consistent with the one
- returned by PyObject_GetAttr. This is allow us to blindly DECREF
- reduce_func at the end of the save() routine.
- */
- Py_INCREF(reduce_func);
Py_INCREF(obj);
reduce_value = _Pickler_FastCall(self, reduce_func, obj);
}
@@ -3359,6 +3368,7 @@
Py_XDECREF(self->output_buffer);
Py_XDECREF(self->write);
Py_XDECREF(self->pers_func);
+ Py_XDECREF(self->dispatch_table);
Py_XDECREF(self->arg);
Py_XDECREF(self->fast_memo);
@@ -3372,6 +3382,7 @@
{
Py_VISIT(self->write);
Py_VISIT(self->pers_func);
+ Py_VISIT(self->dispatch_table);
Py_VISIT(self->arg);
Py_VISIT(self->fast_memo);
return 0;
@@ -3383,6 +3394,7 @@
Py_CLEAR(self->output_buffer);
Py_CLEAR(self->write);
Py_CLEAR(self->pers_func);
+ Py_CLEAR(self->dispatch_table);
Py_CLEAR(self->arg);
Py_CLEAR(self->fast_memo);
@@ -3427,6 +3439,7 @@
PyObject *proto_obj = NULL;
PyObject *fix_imports = Py_True;
_Py_IDENTIFIER(persistent_id);
+ _Py_IDENTIFIER(dispatch_table);
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
kwlist, &file, &proto_obj, &fix_imports))
@@ -3468,6 +3481,13 @@
if (self->pers_func == NULL)
return -1;
}
+ self->dispatch_table = NULL;
+ if (_PyObject_HasAttrId((PyObject *)self, &PyId_dispatch_table)) {
+ self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
+ &PyId_dispatch_table);
+ if (self->dispatch_table == NULL)
+ return -1;
+ }
return 0;
}
@@ -3749,6 +3769,7 @@
static PyMemberDef Pickler_members[] = {
{"bin", T_INT, offsetof(PicklerObject, bin)},
{"fast", T_INT, offsetof(PicklerObject, fast)},
+ {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)},
{NULL}
};
--
Repository URL: http://hg.python.org/cpython
More information about the Python-checkins
mailing list