Project

General

Profile

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