#define PY_SSIZE_T_CLEAN

#include "Python.h"

#include "structmember.h" // PyMemberDef

#define STATIC_CHILDREN 4

#if defined(_MSC_VER)

#define LOCAL(type) static __inline type __fastcall

#define LOCAL(type) static type

#define JOIN_GET(p) ((uintptr_t) (p) & 1)

#define JOIN_SET(p, flag) ((void*) ((uintptr_t) (JOIN_OBJ(p)) | (flag)))

#define JOIN_OBJ(p) ((PyObject*) ((uintptr_t) (p) & ~(uintptr_t)1))

Py_LOCAL_INLINE(void)

_set_joined_ptr(PyObject **p, PyObject *new_joined_ptr)

{

PyObject *tmp = JOIN_OBJ(*p);

*p = new_joined_ptr;

Py_DECREF(tmp);

}

static void _clear_joined_ptr(PyObject **p)

{

if (*p) {

_set_joined_ptr(p, NULL);

}

}

typedef struct {

PyObject *parseerror_obj;

PyObject *deepcopy_obj;

PyObject *elementpath_obj;

PyObject *comment_factory;

PyObject *pi_factory;

/* Interned strings */

PyObject *str_text;

PyObject *str_tail;

PyObject *str_append;

PyObject *str_find;

PyObject *str_findtext;

PyObject *str_findall;

PyObject *str_iterfind;

PyObject *str_doctype;

/* Types defined by this extension */

PyTypeObject *Element_Type;

PyTypeObject *ElementIter_Type;

PyTypeObject *TreeBuilder_Type;

PyTypeObject *XMLParser_Type;

} elementtreestate;

static struct PyModuleDef elementtreemodule;

static inline elementtreestate*

get_elementtree_state(PyObject *module)

{

void *state = PyModule_GetState(module);

assert(state != NULL);

return (elementtreestate *)state;

}

static inline elementtreestate *

get_elementtree_state_by_cls(PyTypeObject *cls)

{

void *state = PyType_GetModuleState(cls);

assert(state != NULL);

return (elementtreestate *)state;

}

static inline elementtreestate *

get_elementtree_state_by_type(PyTypeObject *tp)

{

PyObject *mod = PyType_GetModuleByDef(tp, &elementtreemodule);

assert(mod != NULL);

return get_elementtree_state(mod);

}

static int

elementtree_clear(PyObject *m)

{

elementtreestate *st = get_elementtree_state(m);

Py_CLEAR(st->parseerror_obj);

Py_CLEAR(st->deepcopy_obj);

Py_CLEAR(st->elementpath_obj);

Py_CLEAR(st->comment_factory);

Py_CLEAR(st->pi_factory);

// Interned strings

Py_CLEAR(st->str_append);

Py_CLEAR(st->str_find);

Py_CLEAR(st->str_findall);

Py_CLEAR(st->str_findtext);

Py_CLEAR(st->str_iterfind);

Py_CLEAR(st->str_tail);

Py_CLEAR(st->str_text);

Py_CLEAR(st->str_doctype);

// Heap types

Py_CLEAR(st->Element_Type);

Py_CLEAR(st->ElementIter_Type);

Py_CLEAR(st->TreeBuilder_Type);

Py_CLEAR(st->XMLParser_Type);

return 0;

}

static int

elementtree_traverse(PyObject *m, visitproc visit, void *arg)

{

elementtreestate *st = get_elementtree_state(m);

Py_VISIT(st->parseerror_obj);

Py_VISIT(st->deepcopy_obj);

Py_VISIT(st->elementpath_obj);

Py_VISIT(st->comment_factory);

Py_VISIT(st->pi_factory);

// Heap types

Py_VISIT(st->Element_Type);

Py_VISIT(st->ElementIter_Type);

Py_VISIT(st->TreeBuilder_Type);

Py_VISIT(st->XMLParser_Type);

return 0;

}

static void

elementtree_free(void *m)

{

elementtree_clear((PyObject *)m);

}

LOCAL(PyObject*)

list_join(PyObject* list)

{

/* join list elements */

PyObject* joiner;

PyObject* result;

joiner = PyUnicode_FromStringAndSize("", 0);

if (!joiner)

return NULL;

result = PyUnicode_Join(joiner, list);

Py_DECREF(joiner);

return result;

}

static int

is_empty_dict(PyObject *obj)

{

return PyDict_CheckExact(obj) && PyDict_GET_SIZE(obj) == 0;

}

typedef struct {

/* attributes (a dictionary object), or NULL if no attributes */

PyObject* attrib;

/* child elements */

Py_ssize_t length; /* actual number of items */

Py_ssize_t allocated; /* allocated items */

/* this either points to _children or to a malloced buffer */

PyObject* *children;

PyObject* _children[STATIC_CHILDREN];

} ElementObjectExtra;

typedef struct {

PyObject_HEAD

/* element tag (a string). */

PyObject* tag;

/* text before first child. note that this is a tagged pointer;

PyObject* text;

/* text after this element, in parent. note that this is a tagged

PyObject* tail;

ElementObjectExtra* extra;

PyObject *weakreflist; /* For tp_weaklistoffset */

} ElementObject;

#define Element_CheckExact(st, op) Py_IS_TYPE(op, (st)->Element_Type)

#define Element_Check(st, op) PyObject_TypeCheck(op, (st)->Element_Type)

LOCAL(int)

create_extra(ElementObject* self, PyObject* attrib)

{

self->extra = PyObject_Malloc(sizeof(ElementObjectExtra));

if (!self->extra) {

PyErr_NoMemory();

return -1;

}

self->extra->attrib = Py_XNewRef(attrib);

self->extra->length = 0;

self->extra->allocated = STATIC_CHILDREN;

self->extra->children = self->extra->_children;

return 0;

}

LOCAL(void)

dealloc_extra(ElementObjectExtra *extra)

{

Py_ssize_t i;

if (!extra)

return;

Py_XDECREF(extra->attrib);

for (i = 0; i < extra->length; i++)

Py_DECREF(extra->children[i]);

if (extra->children != extra->_children)

PyObject_Free(extra->children);

PyObject_Free(extra);

}

LOCAL(void)

clear_extra(ElementObject* self)

{

ElementObjectExtra *myextra;

if (!self->extra)

return;

/* Avoid DECREFs calling into this code again (cycles, etc.)

myextra = self->extra;

self->extra = NULL;

dealloc_extra(myextra);

}

LOCAL(PyObject*)

create_new_element(elementtreestate *st, PyObject *tag, PyObject *attrib)

{

ElementObject* self;

self = PyObject_GC_New(ElementObject, st->Element_Type);

if (self == NULL)

return NULL;

self->extra = NULL;

self->tag = Py_NewRef(tag);

self->text = Py_NewRef(Py_None);

self->tail = Py_NewRef(Py_None);

self->weakreflist = NULL;

PyObject_GC_Track(self);

if (attrib != NULL && !is_empty_dict(attrib)) {

if (create_extra(self, attrib) < 0) {

Py_DECREF(self);

return NULL;

}

}

return (PyObject*) self;

}

static PyObject *

element_new(PyTypeObject *type, PyObject *args, PyObject *kwds)

{

ElementObject *e = (ElementObject *)type->tp_alloc(type, 0);

if (e != NULL) {

e->tag = Py_NewRef(Py_None);

e->text = Py_NewRef(Py_None);

e->tail = Py_NewRef(Py_None);

e->extra = NULL;

e->weakreflist = NULL;

}

return (PyObject *)e;

}

static PyObject*

get_attrib_from_keywords(PyObject *kwds)

{

PyObject *attrib_str = PyUnicode_FromString("attrib");

if (attrib_str == NULL) {

return NULL;

}

PyObject *attrib = PyDict_GetItemWithError(kwds, attrib_str);

if (attrib) {

/* If attrib was found in kwds, copy its value and remove it from

if (!PyDict_Check(attrib)) {

Py_DECREF(attrib_str);

PyErr_Format(PyExc_TypeError, "attrib must be dict, not %.100s",

Py_TYPE(attrib)->tp_name);

return NULL;

}

attrib = PyDict_Copy(attrib);

if (attrib && PyDict_DelItem(kwds, attrib_str) < 0) {

Py_SETREF(attrib, NULL);

}

}

else if (!PyErr_Occurred()) {

attrib = PyDict_New();

}

Py_DECREF(attrib_str);

if (attrib != NULL && PyDict_Update(attrib, kwds) < 0) {

Py_DECREF(attrib);

return NULL;

}

return attrib;

}

static int

element_init(PyObject *self, PyObject *args, PyObject *kwds)

{

PyObject *tag;

PyObject *attrib = NULL;

ElementObject *self_elem;

if (!PyArg_ParseTuple(args, "O|O!:Element", &tag, &PyDict_Type, &attrib))

return -1;

if (attrib) {

/* attrib passed as positional arg */

attrib = PyDict_Copy(attrib);

if (!attrib)

return -1;

if (kwds) {

if (PyDict_Update(attrib, kwds) < 0) {

Py_DECREF(attrib);

return -1;

}

}

} else if (kwds) {

/* have keywords args */

attrib = get_attrib_from_keywords(kwds);

if (!attrib)

return -1;

}

self_elem = (ElementObject *)self;

if (attrib != NULL && !is_empty_dict(attrib)) {

if (create_extra(self_elem, attrib) < 0) {

Py_DECREF(attrib);

return -1;

}

}

/* We own a reference to attrib here and it's no longer needed. */

Py_XDECREF(attrib);

/* Replace the objects already pointed to by tag, text and tail. */

Py_XSETREF(self_elem->tag, Py_NewRef(tag));

_set_joined_ptr(&self_elem->text, Py_NewRef(Py_None));

_set_joined_ptr(&self_elem->tail, Py_NewRef(Py_None));

return 0;

}

LOCAL(int)

element_resize(ElementObject* self, Py_ssize_t extra)

{

Py_ssize_t size;

PyObject* *children;

assert(extra >= 0);

/* make sure self->children can hold the given number of extra

if (!self->extra) {

if (create_extra(self, NULL) < 0)

return -1;

}

size = self->extra->length + extra; /* never overflows */

if (size > self->extra->allocated) {

/* use Python 2.4's list growth strategy */

size = (size >> 3) + (size < 9 ? 3 : 6) + size;

/* Coverity CID #182 size_error: Allocating 1 bytes to pointer "children"

size = size ? size : 1;

if ((size_t)size > PY_SSIZE_T_MAX/sizeof(PyObject*))

goto nomemory;

if (self->extra->children != self->extra->_children) {

/* Coverity CID #182 size_error: Allocating 1 bytes to pointer

children = PyObject_Realloc(self->extra->children,

size * sizeof(PyObject*));

if (!children)

goto nomemory;

} else {

children = PyObject_Malloc(size * sizeof(PyObject*));

if (!children)

goto nomemory;

/* copy existing children from static area to malloc buffer */

memcpy(children, self->extra->children,

self->extra->length * sizeof(PyObject*));

}

self->extra->children = children;

self->extra->allocated = size;

}

return 0;

nomemory:

PyErr_NoMemory();

return -1;

}

LOCAL(void)

raise_type_error(PyObject *element)

{

PyErr_Format(PyExc_TypeError,

"expected an Element, not \"%.200s\"",

Py_TYPE(element)->tp_name);

}

LOCAL(int)

element_add_subelement(elementtreestate *st, ElementObject *self,

PyObject *element)

{

/* add a child element to a parent */

if (!Element_Check(st, element)) {

raise_type_error(element);

return -1;

}

if (element_resize(self, 1) < 0)

return -1;

self->extra->children[self->extra->length] = Py_NewRef(element);

self->extra->length++;

return 0;

}

LOCAL(PyObject*)

element_get_attrib(ElementObject* self)

{

/* return borrowed reference to attrib dictionary */

/* note: this function assumes that the extra section exists */

PyObject* res = self->extra->attrib;

if (!res) {

/* create missing dictionary */

res = self->extra->attrib = PyDict_New();

}

return res;

}

LOCAL(PyObject*)

element_get_text(ElementObject* self)

{

/* return borrowed reference to text attribute */

PyObject *res = self->text;

if (JOIN_GET(res)) {

res = JOIN_OBJ(res);

if (PyList_CheckExact(res)) {

PyObject *tmp = list_join(res);

if (!tmp)

return NULL;

self->text = tmp;

Py_SETREF(res, tmp);

}

}

return res;

}

LOCAL(PyObject*)

element_get_tail(ElementObject* self)

{

/* return borrowed reference to text attribute */

PyObject *res = self->tail;

if (JOIN_GET(res)) {

res = JOIN_OBJ(res);

if (PyList_CheckExact(res)) {

PyObject *tmp = list_join(res);

if (!tmp)

return NULL;

self->tail = tmp;

Py_SETREF(res, tmp);

}

}

return res;

}

static PyObject*

subelement(PyObject *self, PyObject *args, PyObject *kwds)

{

PyObject* elem;

elementtreestate *st = get_elementtree_state(self);

ElementObject* parent;

PyObject* tag;

PyObject* attrib = NULL;

if (!PyArg_ParseTuple(args, "O!O|O!:SubElement",

st->Element_Type, &parent, &tag,

&PyDict_Type, &attrib)) {

return NULL;

}

if (attrib) {

/* attrib passed as positional arg */

attrib = PyDict_Copy(attrib);

if (!attrib)

return NULL;

if (kwds != NULL && PyDict_Update(attrib, kwds) < 0) {

Py_DECREF(attrib);

return NULL;

}

} else if (kwds) {

/* have keyword args */

attrib = get_attrib_from_keywords(kwds);

if (!attrib)

return NULL;

} else {

/* no attrib arg, no kwds, so no attribute */

}

elem = create_new_element(st, tag, attrib);

Py_XDECREF(attrib);

if (elem == NULL)

return NULL;

if (element_add_subelement(st, parent, elem) < 0) {

Py_DECREF(elem);

return NULL;

}

return elem;

}

static int

element_gc_traverse(ElementObject *self, visitproc visit, void *arg)

{

Py_VISIT(Py_TYPE(self));

Py_VISIT(self->tag);

Py_VISIT(JOIN_OBJ(self->text));

Py_VISIT(JOIN_OBJ(self->tail));

if (self->extra) {

Py_ssize_t i;

Py_VISIT(self->extra->attrib);

for (i = 0; i < self->extra->length; ++i)

Py_VISIT(self->extra->children[i]);

}

return 0;

}

static int

element_gc_clear(ElementObject *self)

{

Py_CLEAR(self->tag);

_clear_joined_ptr(&self->text);

_clear_joined_ptr(&self->tail);

/* After dropping all references from extra, it's no longer valid anyway,

clear_extra(self);

return 0;

}

static void

element_dealloc(ElementObject* self)

{

PyTypeObject *tp = Py_TYPE(self);

/* bpo-31095: UnTrack is needed before calling any callbacks */

PyObject_GC_UnTrack(self);

Py_TRASHCAN_BEGIN(self, element_dealloc)

if (self->weakreflist != NULL)

PyObject_ClearWeakRefs((PyObject *) self);

/* element_gc_clear clears all references and deallocates extra

element_gc_clear(self);

tp->tp_free((PyObject *)self);

Py_DECREF(tp);

Py_TRASHCAN_END

}

static PyObject *

_elementtree_Element_append_impl(ElementObject *self, PyTypeObject *cls,

PyObject *subelement)

{

elementtreestate *st = get_elementtree_state_by_cls(cls);

if (element_add_subelement(st, self, subelement) < 0)

return NULL;

Py_RETURN_NONE;

}

static PyObject *

_elementtree_Element_clear_impl(ElementObject *self)

{

clear_extra(self);

_set_joined_ptr(&self->text, Py_NewRef(Py_None));

_set_joined_ptr(&self->tail, Py_NewRef(Py_None));

Py_RETURN_NONE;

}

static PyObject *

_elementtree_Element___copy___impl(ElementObject *self, PyTypeObject *cls)

{

Py_ssize_t i;

ElementObject* element;

elementtreestate *st = get_elementtree_state_by_cls(cls);

element = (ElementObject*) create_new_element(

st, self->tag, self->extra ? self->extra->attrib : NULL);

if (!element)

return NULL;

Py_INCREF(JOIN_OBJ(self->text));

_set_joined_ptr(&element->text, self->text);

Py_INCREF(JOIN_OBJ(self->tail));

_set_joined_ptr(&element->tail, self->tail);

assert(!element->extra || !element->extra->length);

if (self->extra) {

if (element_resize(element, self->extra->length) < 0) {

Py_DECREF(element);

return NULL;

}

for (i = 0; i < self->extra->length; i++) {

element->extra->children[i] = Py_NewRef(self->extra->children[i]);

}

assert(!element->extra->length);

element->extra->length = self->extra->length;

}

return (PyObject*) element;

}

LOCAL(PyObject *) deepcopy(elementtreestate *, PyObject *, PyObject *);

static PyObject *

_elementtree_Element___deepcopy___impl(ElementObject *self, PyObject *memo)

{

Py_ssize_t i;

ElementObject* element;

PyObject* tag;

PyObject* attrib;

PyObject* text;

PyObject* tail;

PyObject* id;

PyTypeObject *tp = Py_TYPE(self);

elementtreestate *st = get_elementtree_state_by_type(tp);

tag = deepcopy(st, self->tag, memo);

if (!tag)

return NULL;

if (self->extra && self->extra->attrib) {

attrib = deepcopy(st, self->extra->attrib, memo);

if (!attrib) {

Py_DECREF(tag);

return NULL;

}

} else {

attrib = NULL;

}

element = (ElementObject*) create_new_element(st, tag, attrib);

Py_DECREF(tag);

Py_XDECREF(attrib);

if (!element)

return NULL;

text = deepcopy(st, JOIN_OBJ(self->text), memo);

if (!text)

goto error;

_set_joined_ptr(&element->text, JOIN_SET(text, JOIN_GET(self->text)));

tail = deepcopy(st, JOIN_OBJ(self->tail), memo);

if (!tail)

goto error;

_set_joined_ptr(&element->tail, JOIN_SET(tail, JOIN_GET(self->tail)));

assert(!element->extra || !element->extra->length);

if (self->extra) {

if (element_resize(element, self->extra->length) < 0)

goto error;

for (i = 0; i < self->extra->length; i++) {

PyObject* child = deepcopy(st, self->extra->children[i], memo);

if (!child || !Element_Check(st, child)) {

if (child) {

raise_type_error(child);

Py_DECREF(child);

}

element->extra->length = i;

goto error;

}

element->extra->children[i] = child;

}

assert(!element->extra->length);

element->extra->length = self->extra->length;

}

/* add object to memo dictionary (so deepcopy won't visit it again) */

id = PyLong_FromSsize_t((uintptr_t) self);

if (!id)

goto error;

i = PyDict_SetItem(memo, id, (PyObject*) element);

Py_DECREF(id);

if (i < 0)

goto error;

return (PyObject*) element;

error:

Py_DECREF(element);

return NULL;

}

LOCAL(PyObject *)

deepcopy(elementtreestate *st, PyObject *object, PyObject *memo)

{

/* do a deep copy of the given object */

PyObject *stack[2];

/* Fast paths */

if (object == Py_None || PyUnicode_CheckExact(object)) {

return Py_NewRef(object);

}

if (Py_REFCNT(object) == 1) {

if (PyDict_CheckExact(object)) {

PyObject *key, *value;

Py_ssize_t pos = 0;

int simple = 1;

while (PyDict_Next(object, &pos, &key, &value)) {

if (!PyUnicode_CheckExact(key) || !PyUnicode_CheckExact(value)) {

simple = 0;

break;

}

}

if (simple)

return PyDict_Copy(object);

/* Fall through to general case */

}

else if (Element_CheckExact(st, object)) {

return _elementtree_Element___deepcopy___impl(

(ElementObject *)object, memo);

}

}

/* General case */

if (!st->deepcopy_obj) {

PyErr_SetString(PyExc_RuntimeError,

"deepcopy helper not found");

return NULL;

}

stack[0] = object;

stack[1] = memo;

return _PyObject_FastCall(st->deepcopy_obj, stack, 2);

}

static size_t

_elementtree_Element___sizeof___impl(ElementObject *self)

{

size_t result = _PyObject_SIZE(Py_TYPE(self));

if (self->extra) {

result += sizeof(ElementObjectExtra);

if (self->extra->children != self->extra->_children) {

result += (size_t)self->extra->allocated * sizeof(PyObject*);

}

}

return result;

}

#define PICKLED_TAG "tag"

#define PICKLED_CHILDREN "_children"

#define PICKLED_ATTRIB "attrib"

#define PICKLED_TAIL "tail"

#define PICKLED_TEXT "text"

static PyObject *

_elementtree_Element___getstate___impl(ElementObject *self)

{

Py_ssize_t i;

PyObject *children, *attrib;

/* Build a list of children. */

children = PyList_New(self->extra ? self->extra->length : 0);

if (!children)

return NULL;

for (i = 0; i < PyList_GET_SIZE(children); i++) {

PyObject *child = Py_NewRef(self->extra->children[i]);

PyList_SET_ITEM(children, i, child);

}

if (self->extra && self->extra->attrib) {

attrib = Py_NewRef(self->extra->attrib);

}

else {

attrib = PyDict_New();

if (!attrib) {

Py_DECREF(children);

return NULL;

}

}

return Py_BuildValue("{sOsNsNsOsO}",

PICKLED_TAG, self->tag,

PICKLED_CHILDREN, children,

PICKLED_ATTRIB, attrib,

PICKLED_TEXT, JOIN_OBJ(self->text),

PICKLED_TAIL, JOIN_OBJ(self->tail));

}

static PyObject *

element_setstate_from_attributes(elementtreestate *st,

PyObject *children)

{