msg_tool\ext/
rcdom.rs

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