add _elementtree - CAPI Implementation of ElementTree

Makefile

@@ -1208,7 +1208,8 @@ SHAREDMODS_SRCS := \
 	expat/xmltok.c \
 	expat/xmltok_impl.c \
 	expat/xmltok_ns.c \
- 	pyexpat.c
+ 	pyexpat.c \
+ 	_elementtree.c
 SHAREDMODS_SRCS := $(SHAREDMODS_SRCS:%=from_cpython/Modules/%)
 SHAREDMODS_OBJS := $(SHAREDMODS_NAMES:%=lib_pyston/%.pyston.so)
 

from_cpython/CMakeLists.txt

@@ -112,7 +112,10 @@ file(GLOB_RECURSE STDPARSER_SRCS Parser
 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wno-missing-field-initializers -Wno-tautological-compare -Wno-type-limits -Wno-unused-result -Wno-strict-aliasing")
 add_library(FROM_CPYTHON OBJECT ${STDMODULE_SRCS} ${STDOBJECT_SRCS} ${STDPYTHON_SRCS} ${STDPARSER_SRCS})
 
-add_custom_command(OUTPUT ${CMAKE_BINARY_DIR}/lib_pyston/_multiprocessing.pyston.so ${CMAKE_BINARY_DIR}/lib_pyston/pyexpat.pyston.so
+add_custom_command(OUTPUT
+                       ${CMAKE_BINARY_DIR}/lib_pyston/_multiprocessing.pyston.so
+                       ${CMAKE_BINARY_DIR}/lib_pyston/pyexpat.pyston.so
+                       ${CMAKE_BINARY_DIR}/lib_pyston/_elementtree.pyston.so
                    COMMAND ${CMAKE_BINARY_DIR}/pyston setup.py build --build-lib ${CMAKE_BINARY_DIR}/lib_pyston
                    DEPENDS
                        pyston
@@ -127,5 +130,6 @@ add_custom_command(OUTPUT ${CMAKE_BINARY_DIR}/lib_pyston/_multiprocessing.pyston
                        Modules/expat/xmltok_impl.c
                        Modules/expat/xmltok_ns.c
                        Modules/pyexpat.c
+                       Modules/_elementtree.c
                    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
 add_custom_target(sharedmods DEPENDS ${CMAKE_BINARY_DIR}/lib_pyston/_multiprocessing.pyston.so)

from_cpython/Modules/_elementtree.c

@@ -1791,8 +1791,9 @@ treebuilder_handle_data(TreeBuilderObject* self, PyObject* data)
         Py_INCREF(data); self->data = data;
     } else {
         /* more than one item; use a list to collect items */
-        if (PyString_CheckExact(self->data) && Py_REFCNT(self->data) == 1 &&
-            PyString_CheckExact(data) && PyString_GET_SIZE(data) == 1) {
+        // Pyston change: Py_REFCNT(self->data) -> 2
+        if (PyString_CheckExact(self->data) && /*Py_REFCNT(self->data)*/2 == 1 &&
+              PyString_CheckExact(data) && PyString_GET_SIZE(data) == 1) {
             /* expat often generates single character data sections; handle
                the most common case by resizing the existing string... */
             Py_ssize_t size = PyString_GET_SIZE(self->data);
@@ -2878,6 +2879,16 @@ init_elementtree(void)
     Py_TYPE(&XMLParser_Type) = &PyType_Type;
 #endif
 
+    // Pyston change: register types
+    if (PyType_Ready(&Element_Type) < 0)
+        return;
+    if (PyType_Ready(&TreeBuilder_Type) < 0)
+        return;
+#if defined(USE_EXPAT)
+    if (PyType_Ready(&XMLParser_Type) < 0)
+        return;
+#endif
+
     m = Py_InitModule("_elementtree", _functions);
     if (!m)
         return;

from_cpython/setup.py

@@ -40,8 +40,21 @@ def pyexpat_ext():
                       depends = expat_depends,
                     )
 
+def elementtree_ext():
+    # elementtree depends on expat
+    pyexpat = pyexpat_ext()
+    define_macros = pyexpat.define_macros + [('USE_PYEXPAT_CAPI', None),]
+    return Extension('_elementtree',
+                        define_macros = define_macros,
+                        include_dirs = pyexpat.include_dirs,
+                        libraries = pyexpat.libraries,
+                        sources = [relpath('Modules/_elementtree.c')],
+                        depends = pyexpat.depends,
+                      )
+
+
 setup(name="Pyston",
         version="1.0",
         description="Pyston shared modules",
-        ext_modules=[multiprocessing_ext(), pyexpat_ext()]
+        ext_modules=[multiprocessing_ext(), pyexpat_ext(), elementtree_ext()]
     )

test/tests/elementree_test.py

+# test is based on the cpython ElementTree doc
+def test(ET):
+    xml_str = """<?xml version="1.0"?>
+                <data>
+                    <country name="Liechtenstein">
+                        <rank>1</rank>
+                        <year>2008</year>
+                        <gdppc>141100</gdppc>
+                        <neighbor name="Austria" direction="E"/>
+                        <neighbor name="Switzerland" direction="W"/>
+                    </country>
+                    <country name="Singapore">
+                        <rank>4</rank>
+                        <year>2011</year>
+                        <gdppc>59900</gdppc>
+                        <neighbor name="Malaysia" direction="N"/>
+                    </country>
+                    <country name="Panama">
+                        <rank>68</rank>
+                        <year>2011</year>
+                        <gdppc>13600</gdppc>
+                        <neighbor name="Costa Rica" direction="W"/>
+                        <neighbor name="Colombia" direction="E"/>
+                    </country>
+                </data>
+              """
+
+    root = ET.fromstring(xml_str)
+
+    for child in root:
+        print child.tag, child.attrib
+    print root[0][1].text
+
+    for neighbor in root.iter('neighbor'):
+        print sorted(neighbor.attrib.items())
+        
+    for country in root.findall('country'):
+        rank = country.find('rank').text
+        name = country.get('name')
+        print name, rank
+
+import xml.etree.ElementTree as Python_ET
+import xml.etree.cElementTree as CAPI_ET
+
+test(Python_ET)
+test(CAPI_ET)