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