Merge pull request #1 from kovidgoyal/LP1595781_adopt_external_docLP1595781_adopt_external_doc

Add tests for external document adoption

2 files changed, 108 insertions, 0 deletions

diff --git a/src/lxml/includes/etree_defs.h b/src/lxml/includes/etree_defs.h
index c95eefe1..f85e71b3 100644
--- a/src/lxml/includes/etree_defs.h
+++ b/src/lxml/includes/etree_defs.h
@@ -267,6 +267,12 @@ static void* lxml_unpack_xmldoc_capsule(PyObject* capsule, int* is_owned) {
     xmlDoc *c_doc;
     void *context;
     *is_owned = 0;
+    if (unlikely_condition(!PyCapsule_IsValid(capsule, (const char*)"libxml2:xmlDoc"))) {
+        PyErr_SetString(
+                PyExc_TypeError,
+                "Not a valid capsule. The capsule argument must be a capsule object with name libxml2:xmlDoc");
+        return NULL;
+    }
     c_doc = (xmlDoc*) PyCapsule_GetPointer(capsule, (const char*)"libxml2:xmlDoc");
     if (unlikely_condition(!c_doc)) return NULL;
 
diff --git a/src/lxml/tests/test_external_document.py b/src/lxml/tests/test_external_document.py
new file mode 100644
index 00000000..65055e4d
--- /dev/null
+++ b/src/lxml/tests/test_external_document.py
@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+"""
+Test cases related to direct loading of external libxml2 documents
+"""
+
+import sys
+import unittest
+
+from common_imports import HelperTestCase, etree
+
+DOC_NAME = b'libxml2:xmlDoc'
+DESTRUCTOR_NAME = b'destructor:xmlFreeDoc'
+
+
+@unittest.skipIf(sys.version_info[:2] < (2, 7),
+                 'Not supported for python < 2.7')
+class ExternalDocumentTestCase(HelperTestCase):
+    def setUp(self):
+        import ctypes
+        from ctypes import pythonapi
+        from ctypes.util import find_library
+
+        def wrap(func, restype, *argtypes):
+            func.restype = restype
+            func.argtypes = list(argtypes)
+            return func
+
+        self.get_capsule_name = wrap(pythonapi.PyCapsule_GetName,
+                                     ctypes.c_char_p, ctypes.py_object)
+        self.capsule_is_valid = wrap(pythonapi.PyCapsule_IsValid, ctypes.c_int,
+                                     ctypes.py_object, ctypes.c_char_p)
+        self.new_capsule = wrap(pythonapi.PyCapsule_New, ctypes.py_object,
+                                ctypes.c_void_p, ctypes.c_char_p,
+                                ctypes.c_void_p)
+        self.set_capsule_name = wrap(pythonapi.PyCapsule_SetName, ctypes.c_int,
+                                     ctypes.py_object, ctypes.c_char_p)
+        self.set_capsule_context = wrap(pythonapi.PyCapsule_SetContext,
+                                        ctypes.c_int, ctypes.py_object,
+                                        ctypes.c_char_p)
+        self.get_capsule_context = wrap(pythonapi.PyCapsule_GetContext,
+                                        ctypes.c_char_p, ctypes.py_object)
+        self.get_capsule_pointer = wrap(pythonapi.PyCapsule_GetPointer,
+                                        ctypes.c_void_p, ctypes.py_object,
+                                        ctypes.c_char_p)
+        self.set_capsule_pointer = wrap(pythonapi.PyCapsule_SetPointer,
+                                        ctypes.c_int, ctypes.py_object,
+                                        ctypes.c_void_p)
+        self.set_capsule_destructor = wrap(pythonapi.PyCapsule_SetDestructor,
+                                           ctypes.c_int, ctypes.py_object,
+                                           ctypes.c_void_p)
+        self.PyCapsule_Destructor = ctypes.CFUNCTYPE(None, ctypes.py_object)
+        libxml2 = ctypes.CDLL(find_library('xml2'))
+        self.create_doc = wrap(libxml2.xmlReadMemory, ctypes.c_void_p,
+                               ctypes.c_char_p, ctypes.c_int, ctypes.c_char_p,
+                               ctypes.c_char_p, ctypes.c_int)
+        self.free_doc = wrap(libxml2.xmlFreeDoc, None, ctypes.c_void_p)
+
+    def as_capsule(self, text, capsule_name=DOC_NAME):
+        if not isinstance(text, bytes):
+            text = text.encode('utf-8')
+        doc = self.create_doc(text, len(text), 'base.xml', 'utf-8', 0)
+        ans = self.new_capsule(doc, capsule_name, None)
+        self.set_capsule_context(ans, DESTRUCTOR_NAME)
+        return ans
+
+    def test_external_document_adoption(self):
+        xml = '<r a="1">t</r>'
+        self.assertRaises(TypeError, etree.adopt_external_document, None)
+        capsule = self.as_capsule(xml)
+        self.assertTrue(self.capsule_is_valid(capsule, DOC_NAME))
+        self.assertEqual(DOC_NAME, self.get_capsule_name(capsule))
+        # Create an lxml tree from the capsule (this is a move not a copy)
+        root = etree.adopt_external_document(capsule).getroot()
+        self.assertIsNone(self.get_capsule_name(capsule))
+        self.assertEqual(root.text, 't')
+        root.text = 'new text'
+        # Now reset the capsule so we can copy it
+        self.assertEqual(0, self.set_capsule_name(capsule, DOC_NAME))
+        self.assertEqual(0, self.set_capsule_context(capsule, b'invalid'))
+        # Create an lxml tree from the capsule (this is a copy not a move)
+        root2 = etree.adopt_external_document(capsule).getroot()
+        self.assertEqual(self.get_capsule_context(capsule), b'invalid')
+        # Check that the modification to the tree using the transferred
+        # document was successful
+        self.assertEqual(root.text, root2.text)
+        # Check that further modifications do not show up in the copy (they are
+        # disjoint)
+        root.text = 'other text'
+        self.assertNotEqual(root.text, root2.text)
+        # delete root and ensure root2 survives
+        del root
+        self.assertEqual(root2.text, 'new text')
+
+
+def test_suite():
+    suite = unittest.TestSuite()
+    suite.addTests([unittest.makeSuite(ExternalDocumentTestCase)])
+    return suite
+
+
+if __name__ == '__main__':
+    print('to test use test.py %s' % __file__)