Project

General

Profile

1
# XML DOM tree manipulation
2

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

    
9
import strings
10
import util
11

    
12
##### Escaping input
13

    
14
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
##### Names
21

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

    
27
##### Element node contents
28

    
29
def is_elem(node): return node.nodeType == Node.ELEMENT_NODE
30

    
31
def is_completely_empty(node): return node.firstChild == None
32

    
33
def has_one_child(node):
34
    return node.firstChild != None and node.firstChild.nextSibling == None
35

    
36
def is_simple(node):
37
    '''Whether every child recursively has no more than one child'''
38
    return (not is_elem(node) or is_completely_empty(node)
39
        or (has_one_child(node) and is_simple(node.firstChild)))
40

    
41
class NodeIter:
42
    def __init__(self, node): self.child = node.firstChild
43
    
44
    def __iter__(self): return self
45
    
46
    def curr(self):
47
        if self.child != None: return self.child
48
        raise StopIteration
49
    
50
    def next(self):
51
        child = self.curr()
52
        self.child = self.child.nextSibling
53
        return child
54

    
55
##### Comments
56

    
57
def is_comment(node): return node.nodeType == Node.COMMENT_NODE
58

    
59
def is_empty(node):
60
    for child in NodeIter(node):
61
        if not is_comment(child): return False
62
    return True
63

    
64
def clean_comment(str_):
65
    '''Sanitizes comment node contents. Strips invalid strings.'''
66
    return re.sub(r'-{2,}', r'-', str_) # comments can't contain '--'
67

    
68
def mk_comment(doc, str_): return doc.createComment(clean_comment(str_))
69

    
70
##### Child nodes that are elements
71

    
72
class NodeElemIter:
73
    def __init__(self, node): self.child = node.firstChild
74
    
75
    def __iter__(self): return self
76
    
77
    def curr(self):
78
        while self.child != None:
79
            if is_elem(self.child): return self.child
80
            self.child = self.child.nextSibling
81
        raise StopIteration
82
    
83
    def next(self):
84
        child = self.curr()
85
        self.child = self.child.nextSibling
86
        return child
87

    
88
def first_elem(node): return NodeElemIter(node).next()
89

    
90
def has_elems(node):
91
    try: first_elem(node); return True
92
    except StopIteration: return False
93

    
94
class NodeElemReverseIter:
95
    def __init__(self, node): self.child = node.lastChild
96
    
97
    def __iter__(self): return self
98
    
99
    def curr(self):
100
        while self.child != None:
101
            if is_elem(self.child): return self.child
102
            self.child = self.child.previousSibling
103
        raise StopIteration
104
    
105
    def next(self):
106
        child = self.curr()
107
        self.child = self.child.previousSibling
108
        return child
109

    
110
def last_elem(node): return NodeElemReverseIter(node).next()
111

    
112
##### Parent nodes
113

    
114
def parent(node):
115
    '''Does not treat the document object as the root node's parent, since it's
116
    not a true element node'''
117
    parent_ = node.parentNode
118
    if parent_ != None and is_elem(parent_): return parent_
119
    else: return None
120

    
121
class NodeParentIter:
122
    '''See parent() for special treatment of root node.
123
    Note that the first element returned is the current node, not its parent.'''
124
    def __init__(self, node): self.node = node
125
    
126
    def __iter__(self): return self
127
    
128
    def curr(self):
129
        if self.node != None: return self.node
130
        else: raise StopIteration
131
    
132
    def next(self):
133
        node = self.curr()
134
        self.node = parent(self.node)
135
        return node
136

    
137
##### Element nodes containing text
138

    
139
def is_text_node(node): return node.nodeType == Node.TEXT_NODE
140

    
141
def is_text(node): return has_one_child(node) and is_text_node(node.firstChild)
142

    
143
def value(node):
144
    if is_elem(node):
145
        iter_ = NodeIter(node)
146
        util.skip(iter_, is_comment)
147
        try: return iter_.next().nodeValue
148
        except StopIteration: return None
149
    else: return node.nodeValue
150

    
151
def is_whitespace(node): return is_text_node(node) and value(node).isspace()
152

    
153
def set_value(node, value):
154
    if is_elem(node): node.appendChild(node.ownerDocument.createTextNode(value))
155
    else: node.nodeValue = value
156

    
157
class NodeTextEntryIter:
158
    def __init__(self, node): self.iter_ = NodeElemIter(node)
159
    
160
    def __iter__(self): return self
161
    
162
    def next(self):
163
        util.skip(self.iter_, is_empty)
164
        entry = self.iter_.next()
165
        if is_text(entry): value_ = value(entry)
166
        else:
167
            assert has_one_child(entry) # TODO: convert to an exception
168
            value_ = entry.firstChild
169
        return (entry.tagName, value_)
170

    
171
def is_text_node_entry(val): return util.is_str(val[1])
172

    
173
class TextEntryOnlyIter(util.CheckedIter):
174
    def __init__(self, iterable):
175
        util.CheckedIter.__init__(self, is_text_node_entry, iterable)
176

    
177
##### IDs
178

    
179
def get_id(node):
180
    '''If the node doesn't have an ID, assumes the node itself is the ID.
181
    @return None if the node doesn't have an ID or a value
182
    '''
183
    id_ = node.getAttribute('id')
184
    if id_ != '': return id_
185
    else: return value(node) # assume the node itself is the ID
186

    
187
def set_id(node, id_): node.setAttribute('id', id_)
188

    
189
##### Modifying/replacing a node
190

    
191
def set_child(node, name, value):
192
    '''Note: does not remove any existing child of the same name'''
193
    child = node.ownerDocument.createElement(name)
194
    set_value(child, value)
195
    node.appendChild(child)
196

    
197
def remove(node): node.parentNode.removeChild(node)
198

    
199
def replace(old, new):
200
    '''@param new Node|None'''
201
    if new == None: old.parentNode.removeChild(old)
202
    else: old.parentNode.replaceChild(new, old) # note order reversed
203

    
204
def replace_with_text(node, new):
205
    if util.is_str(new): new = node.ownerDocument.createTextNode(new)
206
    replace(node, new)
207

    
208
##### Searching child nodes
209

    
210
def by_tag_name(node, name, last_only=False, ignore_namespace=False):
211
    '''last_only optimization returns only the last matching node'''
212
    if ignore_namespace: filter_name = strip_namespace
213
    else: filter_name = lambda name: name
214
    name = filter_name(name)
215
    
216
    children = []
217
    if last_only: iter_ = NodeElemReverseIter(node)
218
    else: iter_ = NodeElemIter(node)
219
    for child in iter_:
220
        if filter_name(child.tagName) == name:
221
            children.append(child)
222
            if last_only: break
223
    return children
224

    
225
##### XML documents
226

    
227
def create_doc(root='_'):
228
    return minidom.getDOMImplementation().createDocument(None, root, None)
229

    
230
##### Printing XML
231

    
232
prettyxml_config = dict(addindent='    ', newl='\n')
233
toprettyxml_config = prettyxml_config.copy()
234
util.rename_key(toprettyxml_config, 'addindent', 'indent')
235

    
236
##### minidom modifications
237

    
238
#### Module
239

    
240
minidom._write_data = lambda writer, data: writer.write(escape(data))
241

    
242
minidom.Node.__iter__ = lambda self: NodeIter(self)
243

    
244
def __Node_str(self):
245
    return strings.remove_extra_newl(self.toprettyxml(**toprettyxml_config))
246
minidom.Node.__str__ = __Node_str
247
minidom.Node.__repr__ = __Node_str
248
minidom.Element.__repr__ = __Node_str
249

    
250
#### Node
251

    
252
minidom.Node.pop = lambda self: self.removeChild(self.lastChild)
253

    
254
def __Node_clear(self):
255
    while not is_empty(self): self.pop()
256
minidom.Node.clear = __Node_clear
257

    
258
#### Text
259

    
260
__Text_writexml_orig = minidom.Text.writexml
261
def __Text_writexml(self, *args, **kw_args):
262
    if is_whitespace(self): pass # we add our own whitespace
263
    else: __Text_writexml_orig(self, *args, **kw_args)
264
minidom.Text.writexml = __Text_writexml
265

    
266
#### Attr
267

    
268
def __Attr_str(self): return escape(self.name)+'="'+escape(self.value)+'"'
269
minidom.Attr.__str__ = __Attr_str
270
minidom.Attr.__repr__ = __Attr_str
271

    
272
#### Element
273

    
274
def __Element_write_opening(self, writer, indent='', addindent='', newl=''):
275
    writer.write(indent+'<'+escape(self.tagName))
276
    for attr_idx in xrange(self.attributes.length):
277
        writer.write(' '+str(self.attributes.item(attr_idx)))
278
    writer.write('>'+newl)
279
minidom.Element.write_opening = __Element_write_opening
280

    
281
def __Element_write_closing(self, writer, indent='', addindent='', newl=''):
282
    writer.write('</'+escape(self.tagName)+'>'+newl)
283
minidom.Element.write_closing = __Element_write_closing
284

    
285
__Element_writexml_orig = minidom.Element.writexml
286
def __Element_writexml(self, writer, indent='', addindent='', newl=''):
287
    if isinstance(indent, int): indent = addindent*indent
288
    if is_simple(self):
289
        writer.write(indent)
290
        __Element_writexml_orig(self, writer)
291
        writer.write(newl)
292
    else: __Element_writexml_orig(self, writer, indent, addindent, newl)
293
minidom.Element.writexml = __Element_writexml
294

    
295
#### Document
296

    
297
def __Document_write_opening(self, writer, indent='', addindent='', newl='',
298
    encoding=None):
299
    xmlDecl = '<?xml version="1.0" '
300
    if encoding != None: xmlDecl += 'encoding="'+escape(encoding)+'"'
301
    xmlDecl += '?>'+newl
302
    writer.write(xmlDecl)
303
    assert has_one_child(self)
304
    assert is_elem(self.firstChild)
305
    self.firstChild.write_opening(writer, indent, addindent, newl)
306
minidom.Document.write_opening = __Document_write_opening
307

    
308
def __Document_write_closing(self, writer, indent='', addindent='', newl=''):
309
    self.firstChild.write_closing(writer, indent, addindent, newl)
310
minidom.Document.write_closing = __Document_write_closing
(25-25/29)