use crate::vm::{

AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,

VirtualMachine, atomic_func,

builtins::{PyAnyDict, PyDict, PyDictRef, PyList, PyModule, PyStr, PyType, PyTypeRef},

function::{ArgBytesOrStr, FuncArgs, PySetterValue},

protocol::{PyMappingMethods, PyNumberMethods, PySequenceMethods},

sliceable::{NameElement, SequenceIndex, SliceableSequenceOp},

types::{

AsMapping, AsNumber, AsSequence, Constructor, DefaultConstructor, Initializer,

IterNext, Iterable, Representable, SelfIter,

},

};

use rustpython_common::{lock::PyRwLock, wtf8::Wtf8Buf};

// The per-module bookkeeping `_elementtree.c` keeps in its

// `elementtreestate`: the `Comment`/`ProcessingInstruction` factories

// installed by `_set_factories`, and two helpers imported from Python

// the first time they are needed.

#[pyclass(no_attr, module = "_elementtree", name = "_elementtree_state")]

#[derive(Debug, Default, PyPayload)]

pub(crate) struct ElementTreeState {

comment_factory: PyRwLock<Option<PyObjectRef>>,

pi_factory: PyRwLock<Option<PyObjectRef>>,

element_path: PyRwLock<Option<PyObjectRef>>,

deepcopy: PyRwLock<Option<PyObjectRef>>,

parse_error: PyRwLock<Option<PyTypeRef>>,

}

#[pyclass(flags(DISALLOW_INSTANTIATION))]

impl ElementTreeState {

/// `xml.etree.ElementPath`, imported on first use. `Element.find` and

/// friends delegate every non-trivial path to it, exactly as the C

/// accelerator does.

fn element_path(&self, vm: &VirtualMachine) -> PyResult {

let cached = self.element_path.read().clone();

if let Some(m) = cached {

return Ok(m);

}

// `vm.import()` hands back the *top* package, so the submodule

// itself is picked up from `sys.modules` afterwards.

vm.import("xml.etree.ElementPath", 0)?;

let module = vm

.sys_module

.get_attr("modules", vm)?

.get_item("xml.etree.ElementPath", vm)?;

*self.element_path.write() = Some(module.clone());

Ok(module)

}

fn deepcopy(&self, vm: &VirtualMachine) -> PyResult {

let cached = self.deepcopy.read().clone();

if let Some(f) = cached {

return Ok(f);

}

let f = vm.import("copy", 0)?.get_attr("deepcopy", vm)?;

*self.deepcopy.write() = Some(f.clone());

Ok(f)

}

}

/// The state object stashed on the module by `module_exec`.

///

/// Looked up through `sys.modules` rather than `vm.import()`: the module

/// is always already imported by the time any of its own code runs, and

/// this keeps the lookup down to two dictionary hits.

fn module_state(vm: &VirtualMachine) -> PyResult<PyRef<ElementTreeState>> {

let modules = vm.sys_module.get_attr("modules", vm)?;

let module = match modules.get_item("_elementtree", vm) {

Ok(module) => module,

Err(_) => vm.import("_elementtree", 0)?,

};

module

.get_attr("_state", vm)?

.downcast::<ElementTreeState>()

.map_err(|_| vm.new_runtime_error("_elementtree state is corrupted"))

}

pub(crate) fn module_exec(vm: &VirtualMachine, module: &Py<PyModule>) -> PyResult<()> {

__module_exec(vm, module);

let state = ElementTreeState::default().into_ref(&vm.ctx);

*state.parse_error.write() = module.get_attr("ParseError", vm)?.downcast::<PyType>().ok();

module.set_attr("_state", state, vm)?;

Ok(())

}

#[pyattr(name = "ParseError", once)]

fn parse_error(vm: &VirtualMachine) -> PyTypeRef {

vm.ctx.new_exception_type(

"xml.etree.ElementTree",

"ParseError",

Some(vec![vm.ctx.exceptions.syntax_error.to_owned()]),

)

}

#[pyfunction]

fn _set_factories(

comment_factory: PyObjectRef,

pi_factory: PyObjectRef,

vm: &VirtualMachine,

) -> PyResult<(PyObjectRef, PyObjectRef)> {

if !vm.is_none(&comment_factory) && !comment_factory.is_callable() {

return Err(vm.new_type_error(format!(

"Comment factory must be callable, not {}",

comment_factory.class().name()

)));

}

if !vm.is_none(&pi_factory) && !pi_factory.is_callable() {

return Err(vm.new_type_error(format!(

"PI factory must be callable, not {}",

pi_factory.class().name()

)));

}

let state = module_state(vm)?;

let old_comment = state.comment_factory.read().clone();

let old_pi = state.pi_factory.read().clone();

*state.comment_factory.write() = (!vm.is_none(&comment_factory)).then_some(comment_factory);

*state.pi_factory.write() = (!vm.is_none(&pi_factory)).then_some(pi_factory);

Ok((

old_comment.unwrap_or_else(|| vm.ctx.none()),

old_pi.unwrap_or_else(|| vm.ctx.none()),

))

}

/// `text`/`tail` as the C accelerator stores them: either the value the

/// user sees, or a list of fragments the `TreeBuilder` accumulated that

/// is joined lazily the first time the attribute is read. Keeping the

/// deferred form is what makes character data arriving in many expat

/// callbacks cost one join per element rather than one per callback.

#[derive(Debug)]

struct Joined {

obj: PyObjectRef,

/// `obj` is a list of fragments still to be joined.

pending: bool,

}

impl Joined {

fn plain(obj: PyObjectRef) -> Self {

Self {

obj,

pending: false,

}

}

}

#[derive(Debug)]

struct ElementInner {

tag: PyObjectRef,

text: Joined,

tail: Joined,

/// `None` until an attribute is actually needed, mirroring the C

/// implementation's lazily allocated `extra->attrib`.

attrib: Option<PyRef<PyAnyDict>>,

children: Vec<PyObjectRef>,

}

#[pyattr]

#[pyclass(

#[derive(Debug, PyPayload)]

pub(crate) struct PyElement {

inner: PyRwLock<ElementInner>,

}

// SAFETY: every owned PyObjectRef reachable from an element is visited.

unsafe impl crate::vm::object::Traverse for PyElement {

fn traverse(&self, traverse_fn: &mut crate::vm::object::TraverseFn<'_>) {

let Some(inner) = self.inner.try_read() else {

return;

};

inner.tag.traverse(traverse_fn);

inner.text.obj.traverse(traverse_fn);

inner.tail.obj.traverse(traverse_fn);

inner.attrib.traverse(traverse_fn);

inner.children.traverse(traverse_fn);

}

fn clear(&mut self, out: &mut Vec<PyObjectRef>) {

// Only reached for an object the collector has already proven

// unreachable, so grabbing the lock cannot race with user code.

let Some(mut inner) = self.inner.try_write() else {

return;

};

let none = || -> PyObjectRef { crate::vm::Context::genesis().none.to_owned().into() };

out.append(&mut inner.children);

if let Some(attrib) = inner.attrib.take() {

out.push(attrib.into());

}

out.push(core::mem::replace(&mut inner.tag, none()));

out.push(core::mem::replace(&mut inner.text.obj, none()));

out.push(core::mem::replace(&mut inner.tail.obj, none()));

inner.text.pending = false;

inner.tail.pending = false;

}

}

impl PyElement {

fn new(tag: PyObjectRef, attrib: Option<PyRef<PyAnyDict>>, vm: &VirtualMachine) -> Self {

// CPython elides only empty exact dicts, preserving frozen/subclass identity.

let attrib = attrib.filter(|d| !d.class().is(vm.ctx.types.dict_type) || !d.is_empty());

Self {

inner: PyRwLock::new(ElementInner {

tag,

text: Joined::plain(vm.ctx.none()),

tail: Joined::plain(vm.ctx.none()),

attrib,

children: Vec::new(),

}),

}

}

fn len(&self) -> usize {

self.inner.read().children.len()

}

fn is_empty(&self) -> bool {

self.len() == 0

}

fn child(&self, index: usize) -> Option<PyObjectRef> {

self.inner.read().children.get(index).cloned()

}

pub(crate) fn tag_obj(&self) -> PyObjectRef {

self.inner.read().tag.clone()

}

pub(crate) fn push(&self, child: PyObjectRef) {

self.inner.write().children.push(child);

}

/// Read `text`/`tail`, joining a pending fragment list in place.

fn joined_get(&self, tail: bool, vm: &VirtualMachine) -> PyResult {

{

let inner = self.inner.read();

let slot = if tail { &inner.tail } else { &inner.text };

if !slot.pending {

return Ok(slot.obj.clone());

}

if !slot.obj.downcastable::<PyList>() {

return Ok(slot.obj.clone());

}

}

// The join calls back into Python (str.join), so the lock must be

// released around it and the slot re-checked afterwards.

let list = {

let inner = self.inner.read();

let slot = if tail { &inner.tail } else { &inner.text };

slot.obj.clone()

};

let joined = list_join(&list, vm)?;

let mut inner = self.inner.write();

let slot = if tail {

&mut inner.tail

} else {

&mut inner.text

};

if slot.pending && slot.obj.is(&list) {

*slot = Joined::plain(joined.clone());

}

Ok(joined)

}

fn attrib_or_new(&self, vm: &VirtualMachine) -> PyRef<PyAnyDict> {

let mut inner = self.inner.write();

inner

.attrib

.get_or_insert_with(|| vm.ctx.new_dict().into())

.clone()

}

fn attrib_opt(&self) -> Option<PyRef<PyAnyDict>> {

self.inner.read().attrib.clone()

}

}

/// `"".join(list)`, the C helper `list_join`.

fn list_join(list: &PyObject, vm: &VirtualMachine) -> PyResult {

let empty = vm.ctx.new_str("");

vm.call_method(empty.as_object(), "join", (list.to_owned(),))

}

fn check_element(obj: &PyObject, vm: &VirtualMachine) -> PyResult<()> {

if obj.downcastable::<PyElement>() {

Ok(())

} else {

reject_element(obj, vm)

}

}

fn reject_element(obj: &PyObject, vm: &VirtualMachine) -> PyResult<()> {

Err(vm.new_type_error(format!(

"expected an Element, not \"{}\"",

obj.class().name()

)))

}

fn as_element<'a>(obj: &'a PyObject, vm: &VirtualMachine) -> PyResult<&'a Py<PyElement>> {

obj.downcast_ref::<PyElement>().ok_or_else(|| {

vm.new_type_error(format!(

"expected an Element, not \"{}\"",

obj.class().name()

))

})

}

/// Split `Element(tag, attrib, **extra)` style arguments into the tag and

/// a freshly owned attribute dictionary, matching

/// `get_attrib_from_keywords`: a positional `attrib` is copied, an

/// `attrib` keyword is popped and copied, and the rest of the keywords

/// are merged over it.

fn parse_attrib_args(

args: FuncArgs,

func_name: &str,

extra_positional: usize,

vm: &VirtualMachine,

) -> PyResult<(PyObjectRef, Option<PyDictRef>)> {

let mut args = args;

let nargs = args.args.len();

if nargs < extra_positional + 1 {

return Err(vm.new_type_error(format!(

"{func_name}() takes at least {} argument{} ({nargs} given)",

extra_positional + 1,

if extra_positional == 0 { "" } else { "s" }

)));

}

if nargs > extra_positional + 2 {

return Err(vm.new_type_error(format!(

"{func_name}() takes at most {} arguments ({nargs} given)",

extra_positional + 2

)));

}

let tag = args.args[extra_positional].clone();

let positional_attrib = args.args.get(extra_positional + 1).cloned();

let attrib = match positional_attrib {

Some(obj) => {

let expected = if extra_positional == 0 {

"dict or frozendict"

} else {

"dict"

};

let dict = PyAnyDict::from_object(&obj)

.filter(|_| extra_positional == 0 || obj.downcastable::<PyDict>())

.ok_or_else(|| {

vm.new_type_error(format!(

"{func_name}() argument {} must be {expected}, not {}",

extra_positional + 2,

obj.class().name()

))

})?;

let copied = dict.copy_as_dict(vm)?.into_ref(&vm.ctx);

for (key, value) in args.kwargs.drain(..) {

copied.set_item(&key, value, vm)?;

}

Some(copied)

}

None if !args.kwargs.is_empty() => {

let popped = args.kwargs.shift_remove("attrib");

let copied = match popped {

Some(obj) => {

let dict = PyAnyDict::from_object(&obj).ok_or_else(|| {

vm.new_type_error(format!(

"attrib must be dict or frozendict, not {}",

obj.class().name()

))

})?;

dict.copy_as_dict(vm)?.into_ref(&vm.ctx)

}

None => vm.ctx.new_dict(),

};

for (key, value) in args.kwargs.drain(..) {

copied.set_item(&key, value, vm)?;

}

Some(copied)

}

None => None,

};

Ok((tag, attrib))

}

#[derive(FromArgs)]

struct FindArgs {

#[pyarg(any)]

path: PyObjectRef,

#[pyarg(any, optional)]

namespaces: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct FindTextArgs {

#[pyarg(any)]

path: PyObjectRef,

#[pyarg(any, optional)]

default: Option<PyObjectRef>,

#[pyarg(any, optional)]

namespaces: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct GetArgs {

#[pyarg(any)]

key: PyObjectRef,

#[pyarg(any, optional)]

default: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct IterArgs {

#[pyarg(any, optional)]

tag: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct PiArgs {

#[pyarg(positional)]

target: PyObjectRef,

#[pyarg(positional, optional)]

text: Option<PyObjectRef>,

}

#[derive(FromArgs)]

struct SetEventsArgs {

#[pyarg(positional)]

events_queue: PyObjectRef,

#[pyarg(positional, optional)]

events_to_report: Option<PyObjectRef>,

}

impl Constructor for PyElement {

type Args = FuncArgs;

fn py_new(_cls: &Py<PyType>, _args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {

Ok(Self::new(vm.ctx.none(), None, vm))

}

}

impl Initializer for PyElement {

type Args = FuncArgs;

fn init(zelf: &Py<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {

let (tag, attrib) = parse_attrib_args(args, "Element", 0, vm)?;

let attrib = attrib.filter(|d| !d.is_empty()).map(Into::into);

let _recycle = {

let mut inner = zelf.inner.write();

let (old_attrib, old_children) = if let Some(attrib) = attrib {

(

inner.attrib.replace(attrib),

core::mem::take(&mut inner.children),

)

} else {

(None, Vec::new())

};

(

core::mem::replace(&mut inner.tag, tag),

core::mem::replace(&mut inner.text, Joined::plain(vm.ctx.none())),

core::mem::replace(&mut inner.tail, Joined::plain(vm.ctx.none())),

old_attrib,

old_children,

)

};

Ok(())

}

}

#[pyclass(

impl PyElement {

#[pygetset]

fn tag(zelf: &Py<Self>) -> PyObjectRef {

zelf.inner.read().tag.clone()

}

#[pygetset(setter)]

fn set_tag(zelf: &Py<Self>, value: PySetterValue, vm: &VirtualMachine) -> PyResult<()> {

let value = setter_value(value, vm)?;

// Whatever is replaced may run __del__ on drop, and that code is

// free to touch this very element, so it must not be dropped

// while the lock is held (test_bpo_31728).

let _recycle = core::mem::replace(&mut zelf.inner.write().tag, value);

Ok(())

}

#[pygetset]

fn text(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult {

zelf.joined_get(false, vm)

}

#[pygetset(setter)]

fn set_text(zelf: &Py<Self>, value: PySetterValue, vm: &VirtualMachine) -> PyResult<()> {

let value = setter_value(value, vm)?;

let _recycle = core::mem::replace(&mut zelf.inner.write().text, Joined::plain(value));

Ok(())

}

#[pygetset]

fn tail(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult {

zelf.joined_get(true, vm)

}

#[pygetset(setter)]

fn set_tail(zelf: &Py<Self>, value: PySetterValue, vm: &VirtualMachine) -> PyResult<()> {

let value = setter_value(value, vm)?;

let _recycle = core::mem::replace(&mut zelf.inner.write().tail, Joined::plain(value));

Ok(())

}

#[pygetset]

fn attrib(zelf: &Py<Self>, vm: &VirtualMachine) -> PyRef<PyAnyDict> {

zelf.attrib_or_new(vm)

}

#[pygetset(setter)]

fn set_attrib(zelf: &Py<Self>, value: PySetterValue, vm: &VirtualMachine) -> PyResult<()> {

let value = setter_value(value, vm)?;

let dict = PyAnyDict::from_object(&value)

.ok_or_else(|| {

vm.new_type_error(format!(

"attrib must be dict or frozendict, not {}",

value.class().name()

))

})?

.to_owned();

let _recycle = zelf.inner.write().attrib.replace(dict);

Ok(())

}

#[pymethod]

fn append(zelf: &Py<Self>, subelement: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

check_element(&subelement, vm)?;

zelf.inner.write().children.push(subelement);

Ok(())

}

#[pymethod]

fn extend(zelf: &Py<Self>, elements: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

let items: Vec<PyObjectRef> = elements.try_to_value(vm)?;

// Keep the bulk path while committing valid items before an error.

let valid_prefix = items

.iter()

.position(|element| !element.downcastable::<Self>())

.unwrap_or(items.len());

let mut items = items.into_iter();

zelf.inner

.write()

.children

.extend(items.by_ref().take(valid_prefix));

for element in items {

check_element(&element, vm)?;

zelf.inner.write().children.push(element);

}

Ok(())

}

#[pymethod]

fn insert(

zelf: &Py<Self>,

index: isize,

subelement: PyObjectRef,

vm: &VirtualMachine,

) -> PyResult<()> {

check_element(&subelement, vm)?;

let mut inner = zelf.inner.write();

let len = inner.children.len();

let pos = if index < 0 {

let i = index + len as isize;

if i < 0 { 0 } else { i as usize }

} else {

(index as usize).min(len)

};

inner.children.insert(pos, subelement);

Ok(())

}

#[pymethod]

fn remove(zelf: &Py<Self>, subelement: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

if !subelement.downcastable::<Self>() {

let name = if vm.is_none(&subelement) {

"None".into()

} else {

subelement.class().slot_name()

};

let mut end = name.len().min(50);

while !name.is_char_boundary(end) {

end -= 1;

}

return Err(vm.new_type_error(format!(

"remove() argument must be xml.etree.ElementTree.Element, not {}",

&name[..end]

)));

}

// Identity first, then equality, the way `element_remove` does;

// the equality test can run arbitrary code, so the lock is

// dropped for it and the list re-read afterwards.

let mut found = None;

let mut i = 0;

loop {

let child = {

let inner = zelf.inner.read();

match inner.children.get(i) {

Some(child) => child.clone(),

None => break,

}

};

if child.is(&subelement) {

found = Some(i);

break;

}

if vm.bool_eq(&child, &subelement)? {

found = Some(i);

break;

}

i += 1;

}

let Some(i) = found else {

let mut message = subelement.repr(vm)?.as_wtf8().to_owned();

message.push_str(" not in ");

message.push_wtf8(zelf.as_object().repr(vm)?.as_wtf8());

return Err(vm.new_value_error(message));

};

let _recycle = {

let mut inner = zelf.inner.write();

(i < inner.children.len()).then(|| inner.children.remove(i))

};

Ok(())

}

#[pymethod]

fn clear(zelf: &Py<Self>, vm: &VirtualMachine) {

let _recycle = {

let mut inner = zelf.inner.write();

(

inner.attrib.take(),

core::mem::take(&mut inner.children),

core::mem::replace(&mut inner.text, Joined::plain(vm.ctx.none())),

core::mem::replace(&mut inner.tail, Joined::plain(vm.ctx.none())),

)

};

}

#[pymethod]

fn get(zelf: &Py<Self>, args: GetArgs, vm: &VirtualMachine) -> PyResult {

let default = args.default.unwrap_or_else(|| vm.ctx.none());

let Some(attrib) = zelf.attrib_opt() else {

return Ok(default);

};

Ok(attrib.inner_getitem_opt(&*args.key, vm)?.unwrap_or(default))

}

#[pymethod]

fn set(

zelf: &Py<Self>,

key: PyObjectRef,

value: PyObjectRef,

vm: &VirtualMachine,

) -> PyResult<()> {

let attrib = zelf.attrib_or_new(vm);

attrib.set_item(&*key, value, vm)

}

#[pymethod]

fn keys(zelf: &Py<Self>) -> Vec<PyObjectRef> {

let Some(attrib) = zelf.attrib_opt() else {

return vec![];

};

attrib.keys_vec()

}

#[pymethod]

fn items(zelf: &Py<Self>, vm: &VirtualMachine) -> Vec<PyObjectRef> {

let Some(attrib) = zelf.attrib_opt() else {

return vec![];

};

attrib

.items_vec()

.into_iter()

.map(|(k, v)| vm.ctx.new_tuple(vec![k, v]).into())

.collect()

}

#[pymethod]

fn makeelement(

_zelf: &Py<Self>,

tag: PyObjectRef,

attrib: PyObjectRef,

vm: &VirtualMachine,

) -> PyResult {

let attrib = attrib.downcast::<PyDict>().map_err(|obj| {

vm.new_type_error(format!(

"makeelement() argument 2 must be dict, not {}",

obj.class().name()

))

})?;

let attrib = attrib.copy().into_ref(&vm.ctx);

Ok(create_element(tag, Some(attrib.into()), vm))

}

#[pymethod]

fn find(zelf: &Py<Self>, args: FindArgs, vm: &VirtualMachine) -> PyResult {

let FindArgs { path, namespaces } = args;

let namespaces = namespaces.unwrap_or_else(|| vm.ctx.none());

if check_path(&path) || !vm.is_none(&namespaces) {

let state = module_state(vm)?;

let element_path = state.element_path(vm)?;

return vm.call_method(&element_path, "find", (zelf.to_owned(), path, namespaces));

}

let mut i = 0;

while let Some(child) = zelf.child(i) {

let tag = as_element(&child, vm)?.tag_obj();

if vm.bool_eq(&tag, &path)? {

return Ok(child);

}

i += 1;

}

Ok(vm.ctx.none())

}

#[pymethod]

fn findtext(zelf: &Py<Self>, args: FindTextArgs, vm: &VirtualMachine) -> PyResult {

let FindTextArgs {

path,

default,

namespaces,

} = args;

let default = default.unwrap_or_else(|| vm.ctx.none());

let namespaces = namespaces.unwrap_or_else(|| vm.ctx.none());