Project

General

Profile

1
# XML DOM tree manipulation
2

    
3
import cgi
4
from HTMLParser import HTMLParser
5
import itertools
6
import re
7
from xml.dom import Node
8
import xml.dom.minidom as minidom
9

    
10
import lists
11
import strings
12
import util
13

    
14
##### Escaping input
15

    
16
def escape(str_):
17
    return strings.to_unicode(cgi.escape(str_, True)).encode('ascii',
18
        'xmlcharrefreplace')
19

    
20
def unescape(str_): return HTMLParser().unescape(str_)
21

    
22
##### Names
23

    
24
def strip_namespace(name):
25
    namespace, sep, base = name.partition(':')
26
    if sep != '': return base
27
    else: return name
28

    
29
##### Nodes
30

    
31
def is_node(value): return isinstance(value, Node)
32

    
33
##### Replacing a node
34

    
35
def remove(node): node.parentNode.removeChild(node)
36

    
37
def replace(old, new):
38
    '''@param new Node|None'''
39
    assert old.parentNode != None # not removed from parent tree
40
    if new == None: new = []
41
    else: new = lists.mk_seq(new)
42
    
43
    olds_parent = old.parentNode
44
    if not new: olds_parent.removeChild(old)
45
    else: # at least one node
46
        last = new.pop(-1)
47
        olds_parent.replaceChild(last, old) # note arg order reversed
48
        for node in reversed(new):
49
            olds_parent.insertBefore(node, last) # there is no insertAfter()
50

    
51
def bool2str(val):
52
    '''For use with replace_with_text()'''
53
    if val: return '1'
54
    else: return None # remove node
55

    
56
def replace_with_text(node, new):
57
    '''
58
    @return The *new* node
59
    '''
60
    if isinstance(new, bool): new = bool2str(new)
61
    elif isinstance(new, int) or isinstance(new, float): new = str(new)
62
    if util.is_str(new): new = node.ownerDocument.createTextNode(new)
63
    replace(node, new)
64
    return new
65

    
66
##### Element node contents
67

    
68
def is_elem(node): return node.nodeType == Node.ELEMENT_NODE
69

    
70
def is_completely_empty(node): return node.firstChild == None
71

    
72
def has_one_child(node):
73
    return node.firstChild != None and node.firstChild.nextSibling == None
74

    
75
def only_child(node):
76
    if not has_one_child(node):
77
        raise Exception('Must contain only one child:\n'+strings.ustr(node))
78
    return node.firstChild
79

    
80
def is_simple(node):
81
    '''Whether every child recursively has no more than one child'''
82
    return (not is_elem(node) or is_completely_empty(node)
83
        or (has_one_child(node) and is_simple(node.firstChild)))
84

    
85
class NodeIter:
86
    def __init__(self, node): self.child = node.firstChild
87
    
88
    def __iter__(self): return self
89
    
90
    def curr(self):
91
        if self.child != None: return self.child
92
        raise StopIteration
93
    
94
    def next(self):
95
        child = self.curr()
96
        self.child = self.child.nextSibling
97
        return child
98

    
99
##### Comments
100

    
101
def is_comment(node): return node.nodeType == Node.COMMENT_NODE
102

    
103
def is_empty(node):
104
    for child in NodeIter(node):
105
        if not (is_whitespace(child) or is_comment(child)): return False
106
    return True
107

    
108
def clean_comment(str_):
109
    '''Sanitizes comment node contents. Strips invalid strings.'''
110
    return re.sub(r'-{2,}', r'-', str_) # comments can't contain '--'
111

    
112
def mk_comment(doc, str_): return doc.createComment(clean_comment(str_))
113

    
114
##### Child nodes that are elements
115

    
116
class NodeElemIter:
117
    def __init__(self, node): self.child = node.firstChild
118
    
119
    def __iter__(self): return self
120
    
121
    def curr(self):
122
        while self.child != None:
123
            if is_elem(self.child): return self.child
124
            self.child = self.child.nextSibling
125
        raise StopIteration
126
    
127
    def next(self):
128
        child = self.curr()
129
        self.child = self.child.nextSibling
130
        return child
131

    
132
def first_elem(node): return NodeElemIter(node).next()
133

    
134
def has_elems(node):
135
    try: first_elem(node); return True
136
    except StopIteration: return False
137

    
138
class NodeElemReverseIter:
139
    def __init__(self, node): self.child = node.lastChild
140
    
141
    def __iter__(self): return self
142
    
143
    def curr(self):
144
        while self.child != None:
145
            if is_elem(self.child): return self.child
146
            self.child = self.child.previousSibling
147
        raise StopIteration
148
    
149
    def next(self):
150
        child = self.curr()
151
        self.child = self.child.previousSibling
152
        return child
153

    
154
def last_elem(node): return NodeElemReverseIter(node).next()
155

    
156
class NodeEntryIter:
157
    def __init__(self, node): self.iter_ = NodeElemIter(node)
158
    
159
    def __iter__(self): return self
160
    
161
    def next(self):
162
        entry = self.iter_.next()
163
        values = list(NodeIter(entry))
164
        if not values: values = None
165
        elif len(values) == 1: values = values[0]
166
        return (entry.tagName, values)
167

    
168
##### Parent nodes
169

    
170
def parent(node):
171
    '''Does not treat the document object as the root node's parent, since it's
172
    not a true element node'''
173
    parent_ = node.parentNode
174
    if parent_ != None and is_elem(parent_): return parent_
175
    else: return None
176

    
177
class NodeParentIter:
178
    '''See parent() for special treatment of root node.
179
    Note that the first element returned is the current node, not its parent.'''
180
    def __init__(self, node): self.node = node
181
    
182
    def __iter__(self): return self
183
    
184
    def curr(self):
185
        if self.node != None: return self.node
186
        else: raise StopIteration
187
    
188
    def next(self):
189
        node = self.curr()
190
        self.node = parent(self.node)
191
        return node
192

    
193
##### Element nodes containing text
194

    
195
def is_text_node(node): return node.nodeType == Node.TEXT_NODE
196

    
197
def is_text(node): return has_one_child(node) and is_text_node(node.firstChild)
198

    
199
def value_node(node):
200
    if is_elem(node):
201
        iter_ = NodeIter(node)
202
        util.skip(iter_, is_comment)
203
        try: return iter_.next()
204
        except StopIteration: return None
205
    else: return node
206

    
207
def value(node):
208
    value_ = util.do_ignore_none(lambda n: n.nodeValue, value_node(node))
209
    if value_ == strings.isspace_none_str: value_ = None # None equiv
210
    return value_
211

    
212
def is_whitespace(node):
213
    return is_text_node(node) and (node.nodeValue == ''
214
        or node.nodeValue.isspace())
215

    
216
def set_value(node, value):
217
    value_node_ = value_node(node)
218
    if value != None:
219
        if value_node_ != None:
220
            value_node_.nodeValue = value
221
        else:
222
            assert is_elem(node)
223
            node.appendChild(node.ownerDocument.createTextNode(value))
224
    elif value_node_ != None:
225
        if is_elem(node): remove(value_node_)
226
        else:
227
            if is_text_node(node): value = strings.isspace_none_str # None equiv
228
            node.nodeValue = value
229

    
230
class NodeTextEntryIter:
231
    def __init__(self, node): self.iter_ = NodeElemIter(node)
232
    
233
    def __iter__(self): return self
234
    
235
    def next(self):
236
        entry = self.iter_.next()
237
        if is_empty(entry): value_ = None
238
        elif is_text(entry): value_ = value(entry)
239
        else: value_ = only_child(entry)
240
        return (entry.tagName, value_)
241

    
242
def is_text_node_entry(val): return util.is_str(val[1])
243

    
244
def non_empty(iterable):
245
    return itertools.ifilter(lambda i: i[1] != None, iterable)
246

    
247
class TextEntryOnlyIter(util.CheckedIter):
248
    def __init__(self, iterable):
249
        util.CheckedIter.__init__(self, is_text_node_entry, non_empty(iterable))
250

    
251
##### IDs
252

    
253
def get_id(node):
254
    '''If the node doesn't have an ID, assumes the node itself is the ID.
255
    @return None if the node doesn't have an ID or a value
256
    '''
257
    id_ = node.getAttribute('id')
258
    if id_ != '': return id_
259
    else: return value(node) # assume the node itself is the ID
260

    
261
def set_id(node, id_): node.setAttribute('id', id_)
262

    
263
##### Child nodes
264

    
265
def prune_empty(node):
266
    '''Removes node if it's empty'''
267
    if is_empty(node): remove(node)
268

    
269
def prune_parent(node):
270
    '''Removes node and then parent if it's empty'''
271
    parent = node.parentNode # save parent before detaching node
272
    remove(node)
273
    prune_empty(parent)
274

    
275
def prune_children(node):
276
    '''Removes empty children'''
277
    for child in NodeElemIter(node): prune_empty(child)
278

    
279
def prune(node):
280
    '''Removes empty children and then node if it's empty'''
281
    prune_children(node)
282
    prune_empty(node)
283

    
284
def set_child(node, name, value):
285
    '''Note: does not remove any existing child of the same name'''
286
    child = node.ownerDocument.createElement(name)
287
    set_value(child, value)
288
    node.appendChild(child)
289

    
290
def by_tag_name(node, name, last_only=False, ignore_namespace=False):
291
    '''last_only optimization returns only the last matching node'''
292
    if ignore_namespace: filter_name = strip_namespace
293
    else: filter_name = lambda name: name
294
    name = filter_name(name)
295
    
296
    children = []
297
    if last_only: iter_ = NodeElemReverseIter(node)
298
    else: iter_ = NodeElemIter(node)
299
    for child in iter_:
300
        if filter_name(child.tagName) == name:
301
            children.append(child)
302
            if last_only: break
303
    return children
304

    
305
def merge(from_, into):
306
    '''Merges two nodes of the same tag name and their newly-adjacent children.
307
    @post The into node is saved; the from_ node is deleted.
308
    '''
309
    if from_ == None or into == None: return # base case
310
    if from_.tagName != into.tagName: return # not mergeable
311
    
312
    from_first = from_.firstChild # save before merge
313
    for child in NodeIter(from_): into.appendChild(child)
314
    remove(from_)
315
    
316
    # Recurse
317
    merge(from_first, from_first.previousSibling) # = into.lastChild
318

    
319
def merge_by_name(root, name):
320
    '''Merges siblings in root with the given name'''
321
    children = by_tag_name(root, name)
322
    child0 = children.pop(0)
323
    for child in children: merge(child, child0)
324

    
325
##### XML documents
326

    
327
def create_doc(root='_'):
328
    return minidom.getDOMImplementation().createDocument(None, root, None)
329

    
330
##### Printing XML
331

    
332
prettyxml_config = dict(addindent='    ', newl='\n')
333
toprettyxml_config = prettyxml_config.copy()
334
util.rename_key(toprettyxml_config, 'addindent', 'indent')
335

    
336
##### minidom modifications
337

    
338
#### Module
339

    
340
minidom._write_data = lambda writer, data: writer.write(escape(data))
341

    
342
minidom.Node.__iter__ = lambda self: NodeIter(self)
343

    
344
def __Node_str(self):
345
    return strings.remove_extra_newl(self.toprettyxml(**toprettyxml_config))
346
minidom.Node.__str__ = __Node_str
347
minidom.Node.__repr__ = __Node_str
348
minidom.Element.__repr__ = __Node_str
349

    
350
#### Node
351

    
352
minidom.Node.pop = lambda self: self.removeChild(self.lastChild)
353

    
354
def __Node_clear(self):
355
    while not is_empty(self): self.pop()
356
minidom.Node.clear = __Node_clear
357

    
358
#### Text
359

    
360
__Text_writexml_orig = minidom.Text.writexml
361
def __Text_writexml(self, *args, **kw_args):
362
    if is_whitespace(self): pass # we add our own whitespace
363
    else: __Text_writexml_orig(self, *args, **kw_args)
364
minidom.Text.writexml = __Text_writexml
365

    
366
#### Attr
367

    
368
def __Attr_str(self): return escape(self.name)+'="'+escape(self.value)+'"'
369
minidom.Attr.__str__ = __Attr_str
370
minidom.Attr.__repr__ = __Attr_str
371

    
372
#### Element
373

    
374
def __Element_write_opening(self, writer, indent='', addindent='', newl=''):
375
    writer.write(indent+'<'+escape(self.tagName))
376
    for attr_idx in xrange(self.attributes.length):
377
        writer.write(' '+strings.ustr(self.attributes.item(attr_idx)))
378
    writer.write('>'+newl)
379
minidom.Element.write_opening = __Element_write_opening
380

    
381
def __Element_write_closing(self, writer, indent='', addindent='', newl=''):
382
    writer.write('</'+escape(self.tagName)+'>'+newl)
383
minidom.Element.write_closing = __Element_write_closing
384

    
385
__Element_writexml_orig = minidom.Element.writexml
386
def __Element_writexml(self, writer, indent='', addindent='', newl=''):
387
    if isinstance(indent, int): indent = addindent*indent
388
    if is_simple(self):
389
        writer.write(indent)
390
        __Element_writexml_orig(self, writer)
391
        writer.write(newl)
392
    else: __Element_writexml_orig(self, writer, indent, addindent, newl)
393
minidom.Element.writexml = __Element_writexml
394

    
395
#### Document
396

    
397
def __Document_write_opening(self, writer, indent='', addindent='', newl='',
398
    encoding=None):
399
    xmlDecl = '<?xml version="1.0" '
400
    if encoding != None: xmlDecl += 'encoding="'+escape(encoding)+'"'
401
    xmlDecl += '?>'+newl
402
    writer.write(xmlDecl)
403
    assert has_one_child(self)
404
    assert is_elem(self.firstChild)
405
    self.firstChild.write_opening(writer, indent, addindent, newl)
406
minidom.Document.write_opening = __Document_write_opening
407

    
408
def __Document_write_closing(self, writer, indent='', addindent='', newl=''):
409
    self.firstChild.write_closing(writer, indent, addindent, newl)
410
minidom.Document.write_closing = __Document_write_closing
(41-41/45)