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