[svn r4343] Bug #509504: Fixed objectify pickle support, adding support for

ObjectifiedElement subclasses.

--HG--
branch : trunk

2 files changed, 45 insertions, 11 deletions

diff --git a/src/lxml/lxml.objectify.pyx b/src/lxml/lxml.objectify.pyx
index 03a91f68..948635be 100644
--- a/src/lxml/lxml.objectify.pyx
+++ b/src/lxml/lxml.objectify.pyx
@@ -151,6 +151,10 @@ cdef class ObjectifiedElement(ElementBase):
         else:
             return textOf(self._c_node) or u''
 
+    # pickle support for objectified Element
+    def __reduce__(self):
+        return (fromstring, (etree.tostring(self),))
+
     property text:
         def __get__(self):
             return textOf(self._c_node)
@@ -1359,31 +1363,24 @@ cdef object _dump(_Element element, int indent):
 
 
 ################################################################################
-# Pickle support
+# Pickle support for objectified ElementTree
 
 def __unpickleElementTree(data):
     return etree.ElementTree(fromstring(data))
 
-cdef _setupPickle(elementReduceFunction, elementTreeReduceFunction):
+cdef _setupPickle(elementTreeReduceFunction):
     if python.IS_PYTHON3:
         import copyreg
     else:
         import copy_reg as copyreg
-    copyreg.constructor(fromstring)
-    copyreg.constructor(__unpickleElementTree)
-    copyreg.pickle(ObjectifiedElement,
-                   elementReduceFunction, fromstring)
     copyreg.pickle(etree._ElementTree,
                    elementTreeReduceFunction, __unpickleElementTree)
 
-def pickleReduceElement(obj):
-    return (fromstring, (etree.tostring(obj),))
-
 def pickleReduceElementTree(obj):
     return (__unpickleElementTree, (etree.tostring(obj),))
 
-_setupPickle(pickleReduceElement, pickleReduceElementTree)
-del pickleReduceElement, pickleReduceElementTree
+_setupPickle(pickleReduceElementTree)
+del pickleReduceElementTree
 
 ################################################################################
 # Element class lookup
diff --git a/src/lxml/tests/test_objectify.py b/src/lxml/tests/test_objectify.py
index 71d7ce86..c169de91 100644
--- a/src/lxml/tests/test_objectify.py
+++ b/src/lxml/tests/test_objectify.py
@@ -2328,6 +2328,43 @@ class ObjectifyTestCase(HelperTestCase):
             etree.tostring(new_tree),
             etree.tostring(tree))
 
+    def test_pickle_intelement(self):
+        self._test_pickle('<x>42</x>')
+        self._test_pickle(objectify.DataElement(42))
+
+    def test_pickle_floattelement(self):
+        self._test_pickle('<x>42.0</x>')
+        self._test_pickle(objectify.DataElement(42.0))
+
+    def test_pickle_strelement(self):
+        self._test_pickle('<x>Pickle me!</x>')
+        self._test_pickle(objectify.DataElement('Pickle me!'))
+
+    def test_pickle_boolelement(self):
+        self._test_pickle('<x>true</x>')
+        self._test_pickle('<x>false</x>')
+        self._test_pickle(objectify.DataElement(True))
+        self._test_pickle(objectify.DataElement(False))
+
+    def test_pickle_noneelement(self):
+        self._test_pickle('''
+<x xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:nil="true"/>''')
+        self._test_pickle(objectify.DataElement(None))
+
+    def _test_pickle(self, stringOrElt):
+        import pickle
+        if isinstance(stringOrElt, (etree._Element, etree._ElementTree)):
+            elt = stringOrElt
+        else:
+            elt = self.XML(stringOrElt)
+        out = BytesIO()
+        pickle.dump(elt, out)
+
+        new_elt = pickle.loads(out.getvalue())
+        self.assertEquals(
+            etree.tostring(new_elt),
+            etree.tostring(elt))
+
     # E-Factory tests, need to use sub-elements as root element is always
     # type-looked-up as ObjectifiedElement (no annotations)
     def test_efactory_int(self):