msg_tool\ext/
rcdom.rs

1//! Extensions for markup5ever_rcdom crate.
2use crate::utils::html5ever_arcdom::{AtomicAttribute, Node, NodeData};
3use anyhow::Result;
4use std::sync::{Arc, Mutex, Weak};
5
6/// Extensions for [Node]
7pub trait NodeExt {
8    /// Checks if the node is an element with the given name.
9    ///
10    /// This function ignore namespaces.
11    fn is_element<S: AsRef<str> + ?Sized>(&self, name: &S) -> bool;
12    /// Checks if the node is a processing instruction with the given name.
13    fn is_processing_instruction<S: AsRef<str> + ?Sized>(&self, name: &S) -> bool;
14    /// Returns an iterator over the attribute keys of the element node.
15    ///
16    /// This function returns an empty iterator if the node is not an element.
17    /// Only the local names of the attributes are returned, ignoring namespaces.
18    fn element_attr_keys<'a>(&'a self) -> Result<Box<dyn Iterator<Item = String> + 'a>>;
19    /// Gets the value of the attribute with the given name.
20    ///
21    /// This function returns `Ok(None)` if the node is not an element or if the attribute does not exist.
22    /// This function ignores namespaces and only checks the local name of the attribute.
23    fn get_attr_value<S: AsRef<str> + ?Sized>(&self, name: &S) -> Result<Option<String>>;
24    /// Sets the value of the attribute with the given name.
25    ///
26    /// This function creates the attribute if it does not exist.
27    /// This function ignores namespaces and only checks the local name of the attribute.
28    /// Returns `Ok(())` if the node is not an element.
29    fn set_attr_value<S: AsRef<str> + ?Sized, V: AsRef<str> + ?Sized>(
30        &self,
31        name: &S,
32        value: &V,
33    ) -> Result<()>;
34}
35
36/// Extensions for [Rc<Node>]
37pub trait RcNodeExt {
38    /// Pushes a child node to the current node.
39    fn push_child(&self, child: Arc<Node>) -> Result<()>;
40    /// Create a deep clone
41    fn deep_clone(&self, parent: Option<Weak<Node>>) -> Result<Arc<Node>>;
42    /// Create a deep clone with modification of data.
43    fn deep_clone_with_modify<F: Fn(&mut NodeData) -> Result<()>>(
44        &self,
45        parent: Option<Weak<Node>>,
46        modify: F,
47    ) -> Result<Arc<Node>>;
48    /// Changes a child node at the given index by modifying its data.
49    ///
50    /// Deep clones are needed.
51    fn change_child<F: Fn(&mut NodeData) -> Result<()>>(
52        &self,
53        index: usize,
54        modify: F,
55    ) -> Result<()>;
56}
57
58/// Extensions for [NodeData]
59pub trait NodeDataExt {
60    /// clones the node data.
61    fn clone2(&self) -> Result<NodeData>;
62    /// Sets the content of a processing instruction.
63    ///
64    /// Returns `Ok(())` if the node is not a processing instruction.
65    fn set_processing_instruction_content<S: AsRef<str> + ?Sized>(
66        &mut self,
67        content: &S,
68    ) -> Result<()>;
69}
70
71impl NodeExt for Node {
72    fn is_element<S: AsRef<str> + ?Sized>(&self, name: &S) -> bool {
73        match &self.data {
74            NodeData::Element { name: ename, .. } => ename.local.as_ref() == name.as_ref(),
75            _ => false,
76        }
77    }
78
79    fn is_processing_instruction<S: AsRef<str> + ?Sized>(&self, name: &S) -> bool {
80        match &self.data {
81            NodeData::ProcessingInstruction { target, .. } => target == name.as_ref(),
82            _ => false,
83        }
84    }
85
86    fn element_attr_keys<'a>(&'a self) -> Result<Box<dyn Iterator<Item = String> + 'a>> {
87        match &self.data {
88            NodeData::Element { attrs, .. } => {
89                let attrs = attrs.lock().unwrap();
90                Ok(Box::new(KeyIter { attrs, index: 0 }))
91            }
92            _ => Ok(Box::new(std::iter::empty())),
93        }
94    }
95
96    fn get_attr_value<S: AsRef<str> + ?Sized>(&self, name: &S) -> Result<Option<String>> {
97        match &self.data {
98            NodeData::Element { attrs, .. } => {
99                let attrs = attrs.lock().unwrap();
100                Ok(attrs
101                    .iter()
102                    .find(|a| a.name.local.as_ref() == name.as_ref())
103                    .map(|a| a.value.to_string()))
104            }
105            _ => Ok(None),
106        }
107    }
108
109    fn set_attr_value<S: AsRef<str> + ?Sized, V: AsRef<str> + ?Sized>(
110        &self,
111        name: &S,
112        value: &V,
113    ) -> Result<()> {
114        match &self.data {
115            NodeData::Element { attrs, .. } => {
116                let mut borrowed = attrs.lock().unwrap();
117                if let Some(attr) = borrowed
118                    .iter_mut()
119                    .find(|a| a.name.local.as_ref() == name.as_ref())
120                {
121                    attr.value = value.as_ref().into();
122                } else {
123                    borrowed.push(AtomicAttribute {
124                        name: markup5ever::QualName::new(
125                            None,
126                            markup5ever::Namespace::default(),
127                            name.as_ref().into(),
128                        ),
129                        value: value.as_ref().into(),
130                    });
131                }
132                Ok(())
133            }
134            _ => Ok(()),
135        }
136    }
137}
138
139impl RcNodeExt for Arc<Node> {
140    fn push_child(&self, child: Arc<Node>) -> Result<()> {
141        child.parent.store(Some(Arc::downgrade(self)));
142        self.children.lock().unwrap().push(child);
143        Ok(())
144    }
145
146    fn deep_clone(&self, parent: Option<Weak<Node>>) -> Result<Arc<Node>> {
147        let data = self.data.clone2()?;
148        let node = Node {
149            data,
150            children: Mutex::new(Vec::new()),
151            parent: parent.into(),
152        };
153        let node = Arc::new(node);
154        for child in self.children.lock().unwrap().iter() {
155            let cloned_child = child.deep_clone(Some(Arc::downgrade(&node)))?;
156            node.push_child(cloned_child)?;
157        }
158        Ok(node)
159    }
160
161    fn deep_clone_with_modify<F: Fn(&mut NodeData) -> Result<()>>(
162        &self,
163        parent: Option<Weak<Node>>,
164        modify: F,
165    ) -> Result<Arc<Node>> {
166        let mut data = self.data.clone2()?;
167        modify(&mut data)?;
168        let node = Node {
169            data,
170            children: Mutex::new(Vec::new()),
171            parent: parent.into(),
172        };
173        let node = Arc::new(node);
174        for child in self.children.lock().unwrap().iter() {
175            let cloned_child = child.deep_clone(Some(Arc::downgrade(&node)))?;
176            node.push_child(cloned_child)?;
177        }
178        Ok(node)
179    }
180
181    fn change_child<F: Fn(&mut NodeData) -> Result<()>>(
182        &self,
183        index: usize,
184        modify: F,
185    ) -> Result<()> {
186        let mut children = self.children.lock().unwrap();
187        if index >= children.len() {
188            return Err(anyhow::anyhow!("Index out of bounds"));
189        }
190        let child = children.remove(index);
191        child.parent.take();
192        let nchild = child.deep_clone_with_modify(Some(Arc::downgrade(self)), modify)?;
193        children.insert(index, nchild);
194        Ok(())
195    }
196}
197
198impl NodeDataExt for NodeData {
199    fn clone2(&self) -> Result<Self> {
200        Ok(match self {
201            NodeData::Document => NodeData::Document,
202            NodeData::Comment { contents } => NodeData::Comment {
203                contents: contents.clone(),
204            },
205            NodeData::Doctype {
206                name,
207                public_id,
208                system_id,
209            } => NodeData::Doctype {
210                name: name.clone(),
211                public_id: public_id.clone(),
212                system_id: system_id.clone(),
213            },
214            NodeData::Text { contents } => NodeData::Text {
215                contents: Arc::new(Mutex::new(contents.lock().unwrap().clone())),
216            },
217            NodeData::ProcessingInstruction { target, contents } => {
218                NodeData::ProcessingInstruction {
219                    target: target.clone(),
220                    contents: contents.clone(),
221                }
222            }
223            NodeData::Element {
224                name,
225                attrs,
226                template_contents,
227                mathml_annotation_xml_integration_point,
228            } => {
229                let name = name.clone();
230                let mut nattrs = Vec::new();
231                for attr in attrs.lock().unwrap().iter() {
232                    nattrs.push(AtomicAttribute {
233                        name: attr.name.clone(),
234                        value: attr.value.clone(),
235                    });
236                }
237                let attrs = Mutex::new(nattrs);
238                let template = match template_contents.lock().unwrap().as_ref() {
239                    Some(tc) => Some(tc.deep_clone(None)?),
240                    None => None,
241                };
242                let template_contents = Mutex::new(template);
243                NodeData::Element {
244                    name,
245                    attrs: Arc::new(attrs),
246                    template_contents: Arc::new(template_contents),
247                    mathml_annotation_xml_integration_point:
248                        mathml_annotation_xml_integration_point.clone(),
249                }
250            }
251        })
252    }
253
254    fn set_processing_instruction_content<S: AsRef<str> + ?Sized>(
255        &mut self,
256        content: &S,
257    ) -> Result<()> {
258        match self {
259            NodeData::ProcessingInstruction { contents, .. } => {
260                *contents = content.as_ref().into();
261                Ok(())
262            }
263            _ => Ok(()),
264        }
265    }
266}
267
268struct KeyIter<'a> {
269    attrs: std::sync::MutexGuard<'a, Vec<AtomicAttribute>>,
270    index: usize,
271}
272
273impl<'a> Iterator for KeyIter<'a> {
274    type Item = String;
275
276    fn next(&mut self) -> Option<Self::Item> {
277        if self.index >= self.attrs.len() {
278            None
279        } else {
280            let key = self.attrs[self.index].name.local.as_ref();
281            self.index += 1;
282            Some(key.to_string())
283        }
284    }
285}