Project

General

Profile

1 21 aaronmk
# XML DOM tree manipulation
2
3 73 aaronmk
import cgi
4
from HTMLParser import HTMLParser
5 1809 aaronmk
import re
6 21 aaronmk
from xml.dom import Node
7 299 aaronmk
import xml.dom.minidom as minidom
8 21 aaronmk
9 73 aaronmk
import strings
10 331 aaronmk
import util
11 73 aaronmk
12 840 aaronmk
##### Escaping input
13
14 73 aaronmk
def escape(str_):
15
    return strings.to_unicode(cgi.escape(str_, True)).encode('ascii',
16
        'xmlcharrefreplace')
17
18
def unescape(str_): return HTMLParser().unescape(str_)
19
20 840 aaronmk
##### IDs
21 455 aaronmk
22 21 aaronmk
def get_id(node): return node.getAttribute('id')
23
24
def set_id(node, id_): node.setAttribute('id', id_)
25
26 840 aaronmk
##### Element node contents
27 455 aaronmk
28 840 aaronmk
def is_elem(node): return node.nodeType == Node.ELEMENT_NODE
29
30 453 aaronmk
def is_completely_empty(node): return node.firstChild == None
31 135 aaronmk
32 301 aaronmk
def has_one_child(node):
33
    return node.firstChild != None and node.firstChild.nextSibling == None
34
35 840 aaronmk
def is_simple(node):
36
    '''Whether every child recursively has no more than one child'''
37
    return (not is_elem(node) or is_completely_empty(node)
38
        or (has_one_child(node) and is_simple(node.firstChild)))
39
40 305 aaronmk
class NodeIter:
41
    def __init__(self, node): self.child = node.firstChild
42
43
    def __iter__(self): return self
44
45
    def curr(self):
46
        if self.child != None: return self.child
47
        raise StopIteration
48
49
    def next(self):
50
        child = self.curr()
51
        self.child = self.child.nextSibling
52
        return child
53
54 840 aaronmk
##### Comments
55 455 aaronmk
56 453 aaronmk
def is_comment(node): return node.nodeType == Node.COMMENT_NODE
57
58
def is_empty(node):
59
    for child in NodeIter(node):
60
        if not is_comment(child): return False
61
    return True
62
63 1809 aaronmk
def clean_comment(str_):
64
    '''Sanitizes comment node contents. Strips invalid strings.'''
65
    return re.sub(r'-{2,}', r'-', str_) # comments can't contain '--'
66
67
def mk_comment(doc, str_): return doc.createComment(clean_comment(str_))
68
69 840 aaronmk
##### Child nodes that are elements
70 455 aaronmk
71 21 aaronmk
class NodeElemIter:
72
    def __init__(self, node): self.child = node.firstChild
73
74
    def __iter__(self): return self
75
76
    def curr(self):
77
        while self.child != None:
78 298 aaronmk
            if is_elem(self.child): return self.child
79 21 aaronmk
            self.child = self.child.nextSibling
80
        raise StopIteration
81
82
    def next(self):
83
        child = self.curr()
84
        self.child = self.child.nextSibling
85
        return child
86
87
def first_elem(node): return NodeElemIter(node).next()
88
89 450 aaronmk
def has_elems(node):
90
    try: first_elem(node); return True
91
    except StopIteration: return False
92
93 21 aaronmk
class NodeElemReverseIter:
94
    def __init__(self, node): self.child = node.lastChild
95
96
    def __iter__(self): return self
97
98
    def curr(self):
99
        while self.child != None:
100 298 aaronmk
            if is_elem(self.child): return self.child
101 21 aaronmk
            self.child = self.child.previousSibling
102
        raise StopIteration
103
104
    def next(self):
105
        child = self.curr()
106
        self.child = self.child.previousSibling
107
        return child
108
109
def last_elem(node): return NodeElemReverseIter(node).next()
110
111 840 aaronmk
##### Parent nodes
112
113 1017 aaronmk
def parent(node):
114 999 aaronmk
    '''Does not treat the document object as the root node's parent, since it's
115
    not a true element node'''
116 1017 aaronmk
    parent_ = node.parentNode
117
    if parent_ != None and is_elem(parent_): return parent_
118
    else: return None
119
120
class NodeParentIter:
121
    '''See parent() for special treatment of root node.
122
    Note that the first element returned is the current node, not its parent.'''
123 21 aaronmk
    def __init__(self, node): self.node = node
124
125
    def __iter__(self): return self
126
127
    def curr(self):
128 1017 aaronmk
        if self.node != None: return self.node
129
        else: raise StopIteration
130 21 aaronmk
131
    def next(self):
132
        node = self.curr()
133 1017 aaronmk
        self.node = parent(self.node)
134 21 aaronmk
        return node
135
136 840 aaronmk
##### Element nodes containing text
137 661 aaronmk
138 298 aaronmk
def is_text_node(node): return node.nodeType == Node.TEXT_NODE
139
140 301 aaronmk
def is_text(node): return has_one_child(node) and is_text_node(node.firstChild)
141 21 aaronmk
142
def value(node):
143 29 aaronmk
    if node.firstChild != None: return node.firstChild.nodeValue
144 21 aaronmk
    else: return node.nodeValue
145
146 1720 aaronmk
def is_whitespace(node): return is_text_node(node) and value(node).isspace()
147
148 143 aaronmk
def set_value(node, value):
149 298 aaronmk
    if is_elem(node): node.appendChild(node.ownerDocument.createTextNode(value))
150 22 aaronmk
    else: node.nodeValue = value
151
152 86 aaronmk
class NodeTextEntryIter:
153
    def __init__(self, node): self.iter_ = NodeElemIter(node)
154
155
    def __iter__(self): return self
156
157 1176 aaronmk
    def next(self):
158
        util.skip(self.iter_, is_empty)
159
        entry = self.iter_.next()
160 839 aaronmk
        if is_text(entry): value_ = value(entry)
161
        else:
162
            assert has_one_child(entry) # TODO: convert to an exception
163
            value_ = entry.firstChild
164
        return (entry.tagName, value_)
165 86 aaronmk
166 792 aaronmk
def is_text_node_entry(val): return util.is_str(val[1])
167
168
class TextEntryOnlyIter(util.CheckedIter):
169
    def __init__(self, iterable):
170
        util.CheckedIter.__init__(self, is_text_node_entry, iterable)
171
172 840 aaronmk
##### Modifying/replacing a node
173 455 aaronmk
174 135 aaronmk
def set_child(node, name, value):
175
    '''Note: does not remove any existing child of the same name'''
176
    child = node.ownerDocument.createElement(name)
177 143 aaronmk
    set_value(child, value)
178 135 aaronmk
    node.appendChild(child)
179
180 435 aaronmk
def remove(node): node.parentNode.removeChild(node)
181
182 963 aaronmk
def replace(old, new):
183
    '''@param new Node|None'''
184
    if new == None: old.parentNode.removeChild(old)
185
    else: old.parentNode.replaceChild(new, old) # note order reversed
186 86 aaronmk
187 757 aaronmk
def replace_with_text(node, new):
188 792 aaronmk
    if util.is_str(new): new = node.ownerDocument.createTextNode(new)
189 757 aaronmk
    replace(node, new)
190 86 aaronmk
191 840 aaronmk
##### Searching child nodes
192 455 aaronmk
193 22 aaronmk
def by_tag_name(node, name, last_only=False):
194 135 aaronmk
    '''last_only optimization returns last matching node'''
195 22 aaronmk
    children = []
196 888 aaronmk
    if last_only: iter_ = NodeElemReverseIter(node)
197
    else: iter_ = NodeElemIter(node)
198
    for child in iter_:
199 22 aaronmk
        if child.tagName == name:
200
            children.append(child)
201
            if last_only: break
202
    return children
203 28 aaronmk
204 840 aaronmk
##### XML documents
205 455 aaronmk
206 133 aaronmk
def create_doc(root='_'):
207 303 aaronmk
    return minidom.getDOMImplementation().createDocument(None, root, None)
208 133 aaronmk
209 840 aaronmk
##### Printing XML
210 455 aaronmk
211 304 aaronmk
prettyxml_config = dict(addindent='    ', newl='\n')
212 331 aaronmk
toprettyxml_config = prettyxml_config.copy()
213
util.rename_key(toprettyxml_config, 'addindent', 'indent')
214 304 aaronmk
215 840 aaronmk
##### minidom modifications
216 455 aaronmk
217 1720 aaronmk
#### Module
218
219 299 aaronmk
minidom._write_data = lambda writer, data: writer.write(escape(data))
220 73 aaronmk
221 305 aaronmk
minidom.Node.__iter__ = lambda self: NodeIter(self)
222
223 857 aaronmk
def __Node_str(self):
224 860 aaronmk
    return strings.remove_extra_newl(self.toprettyxml(**toprettyxml_config))
225 796 aaronmk
minidom.Node.__str__ = __Node_str
226
minidom.Node.__repr__ = __Node_str
227
minidom.Element.__repr__ = __Node_str
228 301 aaronmk
229 1720 aaronmk
#### Node
230 888 aaronmk
231 315 aaronmk
minidom.Node.pop = lambda self: self.removeChild(self.lastChild)
232
233
def __Node_clear(self):
234
    while not is_empty(self): self.pop()
235
minidom.Node.clear = __Node_clear
236
237 1720 aaronmk
#### Text
238
239
__Text_writexml_orig = minidom.Text.writexml
240
def __Text_writexml(self, *args, **kw_args):
241
    if is_whitespace(self): pass # we add our own whitespace
242
    else: __Text_writexml_orig(self, *args, **kw_args)
243
minidom.Text.writexml = __Text_writexml
244
245
#### Attr
246
247
def __Attr_str(self): return escape(self.name)+'="'+escape(self.value)+'"'
248
minidom.Attr.__str__ = __Attr_str
249
minidom.Attr.__repr__ = __Attr_str
250
251
#### Element
252
253 301 aaronmk
def __Element_write_opening(self, writer, indent='', addindent='', newl=''):
254 1563 aaronmk
    writer.write(indent+'<'+escape(self.tagName))
255 298 aaronmk
    for attr_idx in xrange(self.attributes.length):
256 891 aaronmk
        writer.write(' '+str(self.attributes.item(attr_idx)))
257 298 aaronmk
    writer.write('>'+newl)
258 301 aaronmk
minidom.Element.write_opening = __Element_write_opening
259 73 aaronmk
260 301 aaronmk
def __Element_write_closing(self, writer, indent='', addindent='', newl=''):
261 1563 aaronmk
    writer.write('</'+escape(self.tagName)+'>'+newl)
262 301 aaronmk
minidom.Element.write_closing = __Element_write_closing
263 298 aaronmk
264 1720 aaronmk
__Element_writexml_orig = minidom.Element.writexml
265 301 aaronmk
def __Element_writexml(self, writer, indent='', addindent='', newl=''):
266 306 aaronmk
    if isinstance(indent, int): indent = addindent*indent
267 774 aaronmk
    if is_simple(self):
268
        writer.write(indent)
269 1720 aaronmk
        __Element_writexml_orig(self, writer)
270 774 aaronmk
        writer.write(newl)
271 1720 aaronmk
    else: __Element_writexml_orig(self, writer, indent, addindent, newl)
272 301 aaronmk
minidom.Element.writexml = __Element_writexml
273 28 aaronmk
274 1720 aaronmk
#### Document
275
276 301 aaronmk
def __Document_write_opening(self, writer, indent='', addindent='', newl='',
277
    encoding=None):
278
    xmlDecl = '<?xml version="1.0" '
279
    if encoding != None: xmlDecl += 'encoding="'+escape(encoding)+'"'
280
    xmlDecl += '?>'+newl
281
    writer.write(xmlDecl)
282
    assert has_one_child(self)
283
    assert is_elem(self.firstChild)
284
    self.firstChild.write_opening(writer, indent, addindent, newl)
285
minidom.Document.write_opening = __Document_write_opening
286
287
def __Document_write_closing(self, writer, indent='', addindent='', newl=''):
288
    self.firstChild.write_closing(writer, indent, addindent, newl)
289
minidom.Document.write_closing = __Document_write_closing