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