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 1814 aaronmk
##### Names
21
22
def strip_namespace(name):
23
    namespace, sep, base = name.partition(':')
24
    if sep != '': return base
25
    else: return name
26
27 2431 aaronmk
##### Nodes
28
29
def is_node(value): return isinstance(value, Node)
30
31 2008 aaronmk
##### Replacing a node
32
33
def remove(node): node.parentNode.removeChild(node)
34
35
def replace(old, new):
36
    '''@param new Node|None'''
37 3329 aaronmk
    assert old.parentNode != None # not removed from parent tree
38
39 2008 aaronmk
    if new == None: old.parentNode.removeChild(old)
40
    else: old.parentNode.replaceChild(new, old) # note order reversed
41
42
def replace_with_text(node, new):
43 3225 aaronmk
    '''
44
    @return The *new* node
45
    '''
46 2008 aaronmk
    if util.is_str(new): new = node.ownerDocument.createTextNode(new)
47
    replace(node, new)
48 3225 aaronmk
    return new
49 2008 aaronmk
50 840 aaronmk
##### Element node contents
51 455 aaronmk
52 840 aaronmk
def is_elem(node): return node.nodeType == Node.ELEMENT_NODE
53
54 453 aaronmk
def is_completely_empty(node): return node.firstChild == None
55 135 aaronmk
56 301 aaronmk
def has_one_child(node):
57
    return node.firstChild != None and node.firstChild.nextSibling == None
58
59 840 aaronmk
def is_simple(node):
60
    '''Whether every child recursively has no more than one child'''
61
    return (not is_elem(node) or is_completely_empty(node)
62
        or (has_one_child(node) and is_simple(node.firstChild)))
63
64 305 aaronmk
class NodeIter:
65
    def __init__(self, node): self.child = node.firstChild
66
67
    def __iter__(self): return self
68
69
    def curr(self):
70
        if self.child != None: return self.child
71
        raise StopIteration
72
73
    def next(self):
74
        child = self.curr()
75
        self.child = self.child.nextSibling
76
        return child
77
78 840 aaronmk
##### Comments
79 455 aaronmk
80 453 aaronmk
def is_comment(node): return node.nodeType == Node.COMMENT_NODE
81
82
def is_empty(node):
83
    for child in NodeIter(node):
84
        if not is_comment(child): return False
85
    return True
86
87 1809 aaronmk
def clean_comment(str_):
88
    '''Sanitizes comment node contents. Strips invalid strings.'''
89
    return re.sub(r'-{2,}', r'-', str_) # comments can't contain '--'
90
91
def mk_comment(doc, str_): return doc.createComment(clean_comment(str_))
92
93 840 aaronmk
##### Child nodes that are elements
94 455 aaronmk
95 21 aaronmk
class NodeElemIter:
96
    def __init__(self, node): self.child = node.firstChild
97
98
    def __iter__(self): return self
99
100
    def curr(self):
101
        while self.child != None:
102 298 aaronmk
            if is_elem(self.child): return self.child
103 21 aaronmk
            self.child = self.child.nextSibling
104
        raise StopIteration
105
106
    def next(self):
107
        child = self.curr()
108
        self.child = self.child.nextSibling
109
        return child
110
111
def first_elem(node): return NodeElemIter(node).next()
112
113 450 aaronmk
def has_elems(node):
114
    try: first_elem(node); return True
115
    except StopIteration: return False
116
117 21 aaronmk
class NodeElemReverseIter:
118
    def __init__(self, node): self.child = node.lastChild
119
120
    def __iter__(self): return self
121
122
    def curr(self):
123
        while self.child != None:
124 298 aaronmk
            if is_elem(self.child): return self.child
125 21 aaronmk
            self.child = self.child.previousSibling
126
        raise StopIteration
127
128
    def next(self):
129
        child = self.curr()
130
        self.child = self.child.previousSibling
131
        return child
132
133
def last_elem(node): return NodeElemReverseIter(node).next()
134
135 840 aaronmk
##### Parent nodes
136
137 1017 aaronmk
def parent(node):
138 999 aaronmk
    '''Does not treat the document object as the root node's parent, since it's
139
    not a true element node'''
140 1017 aaronmk
    parent_ = node.parentNode
141
    if parent_ != None and is_elem(parent_): return parent_
142
    else: return None
143
144
class NodeParentIter:
145
    '''See parent() for special treatment of root node.
146
    Note that the first element returned is the current node, not its parent.'''
147 21 aaronmk
    def __init__(self, node): self.node = node
148
149
    def __iter__(self): return self
150
151
    def curr(self):
152 1017 aaronmk
        if self.node != None: return self.node
153
        else: raise StopIteration
154 21 aaronmk
155
    def next(self):
156
        node = self.curr()
157 1017 aaronmk
        self.node = parent(self.node)
158 21 aaronmk
        return node
159
160 840 aaronmk
##### Element nodes containing text
161 661 aaronmk
162 298 aaronmk
def is_text_node(node): return node.nodeType == Node.TEXT_NODE
163
164 301 aaronmk
def is_text(node): return has_one_child(node) and is_text_node(node.firstChild)
165 21 aaronmk
166 2008 aaronmk
def value_node(node):
167 1832 aaronmk
    if is_elem(node):
168
        iter_ = NodeIter(node)
169
        util.skip(iter_, is_comment)
170 2008 aaronmk
        try: return iter_.next()
171 1832 aaronmk
        except StopIteration: return None
172 2008 aaronmk
    else: return node
173 21 aaronmk
174 2008 aaronmk
def value(node):
175
    return util.do_ignore_none(lambda n: n.nodeValue, value_node(node))
176
177 1720 aaronmk
def is_whitespace(node): return is_text_node(node) and value(node).isspace()
178
179 143 aaronmk
def set_value(node, value):
180 2008 aaronmk
    value_node_ = value_node(node)
181
    if value != None:
182
        if value_node_ != None:
183
            value_node_.nodeValue = value
184
        else:
185
            assert is_elem(node)
186
            node.appendChild(node.ownerDocument.createTextNode(value))
187 2009 aaronmk
    elif value_node_ != None:
188
        if is_elem(node): remove(value_node_)
189
        else: node.nodeValue = value
190 22 aaronmk
191 86 aaronmk
class NodeTextEntryIter:
192
    def __init__(self, node): self.iter_ = NodeElemIter(node)
193
194
    def __iter__(self): return self
195
196 1176 aaronmk
    def next(self):
197
        util.skip(self.iter_, is_empty)
198
        entry = self.iter_.next()
199 839 aaronmk
        if is_text(entry): value_ = value(entry)
200
        else:
201
            assert has_one_child(entry) # TODO: convert to an exception
202
            value_ = entry.firstChild
203
        return (entry.tagName, value_)
204 86 aaronmk
205 792 aaronmk
def is_text_node_entry(val): return util.is_str(val[1])
206
207
class TextEntryOnlyIter(util.CheckedIter):
208
    def __init__(self, iterable):
209
        util.CheckedIter.__init__(self, is_text_node_entry, iterable)
210
211 1836 aaronmk
##### IDs
212
213
def get_id(node):
214
    '''If the node doesn't have an ID, assumes the node itself is the ID.
215
    @return None if the node doesn't have an ID or a value
216
    '''
217
    id_ = node.getAttribute('id')
218
    if id_ != '': return id_
219
    else: return value(node) # assume the node itself is the ID
220
221
def set_id(node, id_): node.setAttribute('id', id_)
222
223 2008 aaronmk
##### Child nodes
224 455 aaronmk
225 135 aaronmk
def set_child(node, name, value):
226
    '''Note: does not remove any existing child of the same name'''
227
    child = node.ownerDocument.createElement(name)
228 143 aaronmk
    set_value(child, value)
229 135 aaronmk
    node.appendChild(child)
230
231 1814 aaronmk
def by_tag_name(node, name, last_only=False, ignore_namespace=False):
232
    '''last_only optimization returns only the last matching node'''
233
    if ignore_namespace: filter_name = strip_namespace
234
    else: filter_name = lambda name: name
235
    name = filter_name(name)
236
237 22 aaronmk
    children = []
238 888 aaronmk
    if last_only: iter_ = NodeElemReverseIter(node)
239
    else: iter_ = NodeElemIter(node)
240
    for child in iter_:
241 1814 aaronmk
        if filter_name(child.tagName) == name:
242 22 aaronmk
            children.append(child)
243
            if last_only: break
244
    return children
245 28 aaronmk
246 3226 aaronmk
def merge(from_, into):
247
    '''The into node is saved; the from_ node is deleted.'''
248
    for child in NodeIter(from_): into.appendChild(child)
249
    remove(from_)
250
251
def merge_adjacent(node):
252 3325 aaronmk
    '''Repeatedly merges a node with its siblings as long as they or their
253
    newly-adjacent children have the same tag name'''
254 3226 aaronmk
    if node == None: return # base case
255
256
    for node_, next in [(node.previousSibling, node), (node, node.nextSibling)]:
257
        if node_ != None and next != None and node_.tagName == next.tagName:
258
            next_first = next.firstChild # save before merge
259
            merge(next, node_)
260
            merge_adjacent(next_first) # previousSibling is now node_.lastChild
261
262 3325 aaronmk
def merge_same_name(node):
263
    '''Merges nodes with the same name as a node, using merge_adjacent()'''
264
    for node_ in by_tag_name(node.parentNode, node.tagName):
265
        if node_ is node: continue # not self
266
        merge_adjacent(node_)
267
268 840 aaronmk
##### XML documents
269 455 aaronmk
270 133 aaronmk
def create_doc(root='_'):
271 303 aaronmk
    return minidom.getDOMImplementation().createDocument(None, root, None)
272 133 aaronmk
273 840 aaronmk
##### Printing XML
274 455 aaronmk
275 304 aaronmk
prettyxml_config = dict(addindent='    ', newl='\n')
276 331 aaronmk
toprettyxml_config = prettyxml_config.copy()
277
util.rename_key(toprettyxml_config, 'addindent', 'indent')
278 304 aaronmk
279 840 aaronmk
##### minidom modifications
280 455 aaronmk
281 1720 aaronmk
#### Module
282
283 299 aaronmk
minidom._write_data = lambda writer, data: writer.write(escape(data))
284 73 aaronmk
285 305 aaronmk
minidom.Node.__iter__ = lambda self: NodeIter(self)
286
287 857 aaronmk
def __Node_str(self):
288 860 aaronmk
    return strings.remove_extra_newl(self.toprettyxml(**toprettyxml_config))
289 796 aaronmk
minidom.Node.__str__ = __Node_str
290
minidom.Node.__repr__ = __Node_str
291
minidom.Element.__repr__ = __Node_str
292 301 aaronmk
293 1720 aaronmk
#### Node
294 888 aaronmk
295 315 aaronmk
minidom.Node.pop = lambda self: self.removeChild(self.lastChild)
296
297
def __Node_clear(self):
298
    while not is_empty(self): self.pop()
299
minidom.Node.clear = __Node_clear
300
301 1720 aaronmk
#### Text
302
303
__Text_writexml_orig = minidom.Text.writexml
304
def __Text_writexml(self, *args, **kw_args):
305
    if is_whitespace(self): pass # we add our own whitespace
306
    else: __Text_writexml_orig(self, *args, **kw_args)
307
minidom.Text.writexml = __Text_writexml
308
309
#### Attr
310
311
def __Attr_str(self): return escape(self.name)+'="'+escape(self.value)+'"'
312
minidom.Attr.__str__ = __Attr_str
313
minidom.Attr.__repr__ = __Attr_str
314
315
#### Element
316
317 301 aaronmk
def __Element_write_opening(self, writer, indent='', addindent='', newl=''):
318 1563 aaronmk
    writer.write(indent+'<'+escape(self.tagName))
319 298 aaronmk
    for attr_idx in xrange(self.attributes.length):
320 891 aaronmk
        writer.write(' '+str(self.attributes.item(attr_idx)))
321 298 aaronmk
    writer.write('>'+newl)
322 301 aaronmk
minidom.Element.write_opening = __Element_write_opening
323 73 aaronmk
324 301 aaronmk
def __Element_write_closing(self, writer, indent='', addindent='', newl=''):
325 1563 aaronmk
    writer.write('</'+escape(self.tagName)+'>'+newl)
326 301 aaronmk
minidom.Element.write_closing = __Element_write_closing
327 298 aaronmk
328 1720 aaronmk
__Element_writexml_orig = minidom.Element.writexml
329 301 aaronmk
def __Element_writexml(self, writer, indent='', addindent='', newl=''):
330 306 aaronmk
    if isinstance(indent, int): indent = addindent*indent
331 774 aaronmk
    if is_simple(self):
332
        writer.write(indent)
333 1720 aaronmk
        __Element_writexml_orig(self, writer)
334 774 aaronmk
        writer.write(newl)
335 1720 aaronmk
    else: __Element_writexml_orig(self, writer, indent, addindent, newl)
336 301 aaronmk
minidom.Element.writexml = __Element_writexml
337 28 aaronmk
338 1720 aaronmk
#### Document
339
340 301 aaronmk
def __Document_write_opening(self, writer, indent='', addindent='', newl='',
341
    encoding=None):
342
    xmlDecl = '<?xml version="1.0" '
343
    if encoding != None: xmlDecl += 'encoding="'+escape(encoding)+'"'
344
    xmlDecl += '?>'+newl
345
    writer.write(xmlDecl)
346
    assert has_one_child(self)
347
    assert is_elem(self.firstChild)
348
    self.firstChild.write_opening(writer, indent, addindent, newl)
349
minidom.Document.write_opening = __Document_write_opening
350
351
def __Document_write_closing(self, writer, indent='', addindent='', newl=''):
352
    self.firstChild.write_closing(writer, indent, addindent, newl)
353
minidom.Document.write_closing = __Document_write_closing