added a flag to encode/decode flattened data as json
[OpenSRF.git] / src / python / osrf / xml_obj.py
1 import xml.dom.minidom
2 import osrf.json
3 from xml.sax import handler, make_parser, saxutils
4 import urllib, re
5
6 def xml_file_to_object(filename):
7     """Turns the contents of an XML file into a Python object"""
8     doc = xml.dom.minidom.parse(filename)
9     obj = xml_node_to_object(doc.documentElement)
10     doc.unlink()
11     return obj
12
13 def xml_string_to_object(string):
14     """Turns an XML string into a Python object"""
15     doc = xml.dom.minidom.parseString(string)
16     obj = xml_node_to_object(doc.documentElement)
17     doc.unlink()
18     return obj
19
20 def xml_node_to_object(xml_node):
21     """Turns an XML node into a Python object"""
22     obj = {}
23
24     if xml_node.nodeType != xml_node.ELEMENT_NODE:
25         return obj
26
27     done = False
28     node_name = xml_node.nodeName
29
30     for node_child in xml_node.childNodes:
31         if node_child.nodeType == xml_node.ELEMENT_NODE:
32             sub_obj = xml_node_to_object(node_child)
33             __append_child_node(obj, node_name, node_child.nodeName, sub_obj)
34             done = True
35
36     for attr in xml_node.attributes.values():
37         __append_child_node(obj, node_name, attr.name,
38             dict([(attr.name, attr.value)]))
39         
40
41     if not done and len(xml_node.childNodes) > 0:
42         # If the node has no element children, clean up the text 
43         # content and use that as the data
44         text_node = xml_node.childNodes[0] # extract the text node
45         data = unicode(text_node.nodeValue).replace('^\s*','')
46         data = data.replace('\s*$','')
47
48         if node_name in obj:
49             # the current element contains attributes and text
50             obj[node_name]['#text'] = data
51         else:
52             # the current element contains text only
53             obj[node_name] = data
54
55     return obj
56
57
58 def __append_child_node(obj, node_name, child_name, sub_obj):
59     """ If a node has element children, create a new sub-object 
60         for this node, attach an array for each type of child
61         and recursively collect the children data into the array(s) """
62
63     if not obj.has_key(node_name):
64         obj[node_name] = {}
65
66     if not obj[node_name].has_key(child_name):
67         # we've encountered 1 sub-node with node_child's name
68         if child_name in sub_obj:
69             obj[node_name][child_name] = sub_obj[child_name]
70         else:
71             obj[node_name][child_name] = None
72
73     else:
74         if isinstance(obj[node_name][child_name], list):
75             # we already have multiple sub-nodes with node_child's name
76             obj[node_name][child_name].append(sub_obj[child_name])
77
78         else:
79             # we already have 1 sub-node with node_child's name, make 
80             # it a list and append the current node
81             val = obj[node_name][child_name]
82             obj[node_name][child_name] = [ val, sub_obj[child_name] ]
83
84
85
86 class XMLFlattener(handler.ContentHandler):
87     ''' Turns an XML string into a flattened dictionary of properties.
88
89         Example <doc><a><b>text1</b></a><c>text2</c><c>text3</c></doc> becomes
90         {
91             'doc.a.b' : 'text1',
92             'doc.c' : ['text2', 'text3']
93         }
94     '''
95
96     reg = re.compile('^\s*$')
97     class Handler(handler.ContentHandler):
98         def __init__(self):
99             self.result = {}
100             self.elements = []
101             self.use_json = None
102     
103         def startElement(self, name, attrs):
104             self.elements.append(name)
105
106         def characters(self, chars):
107             text = urllib.unquote_plus(chars)
108             if re.match(XMLFlattener.reg, text):
109                 return
110             key = ''
111             for elm in self.elements:
112                 key += elm + '.'
113             key = key[:-1]
114             
115             if key in self.result:
116                 data = self._decode(self.result[key])
117                 if isinstance(data, list):
118                     data.append(text)
119                 else:
120                     data = [data, text]
121                 self.result[key] = self._encode(data)
122             else:
123                 self.result[key] = self._encode(text)
124
125             
126         def endElement(self, name):
127             self.elements.pop()
128
129         def _decode(self, string):
130             if self.use_json:
131                 return osrf.json.to_object(string)
132             return string
133
134         def _encode(self, obj):
135             if self.use_json:
136                 return osrf.json.to_json(obj)
137             return obj
138             
139
140
141     def __init__(self, xml_str, encode_as_json=False):
142         self.xml_str = xml_str
143         self.use_json = encode_as_json
144
145     def parse(self):
146         ''' Parses the XML string and returns the dict of keys/values '''
147         sax_handler = XMLFlattener.Handler()
148         sax_handler.use_json = self.use_json
149         parser = make_parser()
150         parser.setContentHandler(sax_handler)
151         try:
152             import StringIO
153             parser.parse(StringIO.StringIO(self.xml_str))
154         except Exception, e:
155             osrf.log.log_error('Error parsing XML: %s' % unicode(e))
156             raise e
157
158         return sax_handler.result
159
160
161
162