Project

General

Profile

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