Fixed a corner case
[sfa.git] / sfa / util / rspec.py
1 ### $Id$
2 ### $URL$
3
4 import sys
5 import pprint
6 import os
7 import httplib
8 from xml.dom import minidom
9 from types import StringTypes, ListType
10
11 class Rspec:
12
13     def __init__(self, xml = None, xsd = None, NSURL = None):
14         '''
15         Class to manipulate RSpecs.  Reads and parses rspec xml into python dicts
16         and reads python dicts and writes rspec xml
17
18         self.xsd = # Schema.  Can be local or remote file.
19         self.NSURL = # If schema is remote, Name Space URL to query (full path minus filename)
20         self.rootNode = # root of the DOM
21         self.dict = # dict of the RSpec.
22         self.schemaDict = {} # dict of the Schema
23         '''
24  
25         self.xsd = xsd
26         self.rootNode = None
27         self.dict = {}
28         self.schemaDict = {}
29         self.NSURL = NSURL 
30         if xml: 
31             if type(xml) == file:
32                 self.parseFile(xml)
33             if type(xml) == str:
34                 self.parseString(xml)
35             self.dict = self.toDict() 
36         if xsd:
37             self._parseXSD(self.NSURL + self.xsd)
38
39
40     def _getText(self, nodelist):
41         rc = ""
42         for node in nodelist:
43             if node.nodeType == node.TEXT_NODE:
44                 rc = rc + node.data
45         return rc
46   
47     # The rspec is comprised of 2 parts, and 1 reference:
48     # attributes/elements describe individual resources
49     # complexTypes are used to describe a set of attributes/elements
50     # complexTypes can include a reference to other complexTypes.
51   
52   
53     def _getName(self, node):
54         '''Gets name of node. If tag has no name, then return tag's localName'''
55         name = None
56         if not node.nodeName.startswith("#"):
57             if node.localName:
58                 name = node.localName
59             elif node.attributes.has_key("name"):
60                 name = node.attributes.get("name").value
61         return name     
62  
63  
64     # Attribute.  {name : nameofattribute, {items: values})
65     def _attributeDict(self, attributeDom):
66         '''Traverse single attribute node.  Create a dict {attributename : {name: value,}]}'''
67         node = {} # parsed dict
68         for attr in attributeDom.attributes.keys():
69             node[attr] = attributeDom.attributes.get(attr).value
70         return node
71   
72  
73     def appendToDictOrCreate(self, dict, key, value):
74         if (dict.has_key(key)):
75             dict[key].append(value)
76         else:
77             dict[key]=[value]
78         return dict
79
80     def toGenDict(self, nodeDom=None, parentdict=None, siblingdict={}, parent=None):
81         """
82         convert an XML to a nested dict:
83           * Non-terminal nodes (elements with string children and attributes) are simple dictionaries
84           * Terminal nodes (the rest) are nested dictionaries
85         """
86
87         if (not nodeDom):
88             nodeDom=self.rootNode
89
90         curNodeName = nodeDom.localName
91
92         if (nodeDom.hasChildNodes()):
93             childdict={}
94             for attribute in nodeDom.attributes.keys():
95                 siblingdict = self.appendToDictOrCreate(siblingdict, attribute, nodeDom.getAttribute(attribute))
96
97             for child in nodeDom.childNodes[:-1]:
98                 if (child.nodeValue):
99                     siblingdict = self.appendToDictOrCreate(siblingdict, curNodeName, child.nodeValue)
100                 else:
101                     childdict = self.toGenDict(child, None, childdict, curNodeName)
102
103             child = nodeDom.childNodes[-1]
104             if (child.nodeValue):
105                 siblingdict = self.appendToDictOrCreate(siblingdict, curNodeName, child.nodeValue)
106             else:
107                 siblingdict = self.toGenDict(child, siblingdict, childdict, curNodeName)
108         else:
109             siblingdict[curNodeName]=[]
110             
111         if (parentdict is not None):
112             parentdict = self.appendToDictOrCreate(parentdict, parent, siblingdict)
113             return parentdict
114         else:
115             return siblingdict
116
117
118
119     def toDict(self, nodeDom = None):
120         """
121         convert this rspec to a dict and return it.
122         """
123         node = {}
124         if not nodeDom:
125              nodeDom = self.rootNode
126   
127         elementName = nodeDom.nodeName
128         if elementName and not elementName.startswith("#"):
129             # attributes have tags and values.  get {tag: value}, else {type: value}
130             node[elementName] = self._attributeDict(nodeDom)
131             # resolve the child nodes.
132             if nodeDom.hasChildNodes():
133                 for child in nodeDom.childNodes:
134                     childName = self._getName(child)
135                     # skip null children 
136                     if not childName:
137                         continue
138                     # initialize the possible array of children        
139                     if not node[elementName].has_key(childName):
140                         node[elementName][childName] = []
141                     # if child node has text child nodes
142                     # append the children to the array as strings
143                     if child.hasChildNodes() and isinstance(child.childNodes[0], minidom.Text):
144                         for nextchild in child.childNodes:
145                             node[elementName][childName].append(nextchild.data)
146                     # convert element child node to dict
147                     else:       
148                         childdict = self.toDict(child)
149                         for value in childdict.values():
150                             node[elementName][childName].append(value)
151                     #node[childName].append(self.toDict(child))
152         return node
153
154   
155     def toxml(self):
156         """
157         convert this rspec to an xml string and return it.
158         """
159         return self.rootNode.toxml()
160
161   
162     def toprettyxml(self):
163         """
164         print this rspec in xml in a pretty format.
165         """
166         return self.rootNode.toprettyxml()
167
168   
169     def parseFile(self, filename):
170         """
171         read a local xml file and store it as a dom object.
172         """
173         dom = minidom.parse(filename)
174         self.rootNode = dom.childNodes[0]
175
176
177     def parseString(self, xml):
178         """
179         read an xml string and store it as a dom object.
180         """
181         xml = xml.replace('\n', '').replace('\t', '').strip()
182         dom = minidom.parseString(xml)
183         self.rootNode = dom.childNodes[0]
184
185  
186     def _httpGetXSD(self, xsdURI):
187         # split the URI into relevant parts
188         host = xsdURI.split("/")[2]
189         if xsdURI.startswith("https"):
190             conn = httplib.HTTPSConnection(host,
191                 httplib.HTTPSConnection.default_port)
192         elif xsdURI.startswith("http"):
193             conn = httplib.HTTPConnection(host,
194                 httplib.HTTPConnection.default_port)
195         conn.request("GET", xsdURI)
196         # If we can't download the schema, raise an exception
197         r1 = conn.getresponse()
198         if r1.status != 200: 
199             raise Exception
200         return r1.read().replace('\n', '').replace('\t', '').strip() 
201
202
203     def _parseXSD(self, xsdURI):
204         """
205         Download XSD from URL, or if file, read local xsd file and set schemaDict
206         """
207         # Since the schema definiton is a global namespace shared by and agreed upon by
208         # others, this should probably be a URL.  Check for URL, download xsd, parse, or 
209         # if local file, use local file.
210         schemaDom = None
211         if xsdURI.startswith("http"):
212             try: 
213                 schemaDom = minidom.parseString(self._httpGetXSD(xsdURI))
214             except Exception, e:
215                 # logging.debug("%s: web file not found" % xsdURI)
216                 # logging.debug("Using local file %s" % self.xsd")
217                 print e
218                 print "Can't find %s on the web. Continuing." % xsdURI
219         if not schemaDom:
220             if os.path.exists(xsdURI):
221                 # logging.debug("using local copy.")
222                 print "Using local %s" % xsdURI
223                 schemaDom = minidom.parse(xsdURI)
224             else:
225                 raise Exception("Can't find xsd locally")
226         self.schemaDict = self.toDict(schemaDom.childNodes[0])
227
228
229     def dict2dom(self, rdict, include_doc = False):
230         """
231         convert a dict object into a dom object.
232         """
233      
234         def elementNode(tagname, rd):
235             element = minidom.Element(tagname)
236             for key in rd.keys():
237                 if isinstance(rd[key], StringTypes) or isinstance(rd[key], int):
238                     element.setAttribute(key, str(rd[key]))
239                 elif isinstance(rd[key], dict):
240                     child = elementNode(key, rd[key])
241                     element.appendChild(child)
242                 elif isinstance(rd[key], list):
243                     for item in rd[key]:
244                         if isinstance(item, dict):
245                             child = elementNode(key, item)
246                             element.appendChild(child)
247                         elif isinstance(item, StringTypes) or isinstance(item, int):
248                             child = minidom.Element(key)
249                             text = minidom.Text()
250                             text.data = item
251                             child.appendChild(text)
252                             element.appendChild(child) 
253             return element
254         
255         # Minidom does not allow documents to have more then one
256         # child, but elements may have many children. Because of
257         # this, the document's root node will be the first key/value
258         # pair in the dictionary.  
259         node = elementNode(rdict.keys()[0], rdict.values()[0])
260         if include_doc:
261             rootNode = minidom.Document()
262             rootNode.appendChild(node)
263         else:
264             rootNode = node
265         return rootNode
266
267  
268     def parseDict(self, rdict, include_doc = True):
269         """
270         Convert a dictionary into a dom object and store it.
271         """
272         self.rootNode = self.dict2dom(rdict, include_doc)
273  
274  
275     def getDictsByTagName(self, tagname, dom = None):
276         """
277         Search the dom for all elements with the specified tagname
278         and return them as a list of dicts
279         """
280         if not dom:
281             dom = self.rootNode
282         dicts = []
283         doms = dom.getElementsByTagName(tagname)
284         dictlist = [self.toDict(d) for d in doms]
285         for item in dictlist:
286             for value in item.values():
287                 dicts.append(value)
288         return dicts
289
290     def getDictByTagNameValue(self, tagname, value, dom = None):
291         """
292         Search the dom for the first element with the specified tagname
293         and value and return it as a dict.
294         """
295         tempdict = {}
296         if not dom:
297             dom = self.rootNode
298         dicts = self.getDictsByTagName(tagname, dom)
299         
300         for rdict in dicts:
301             if rdict.has_key('name') and rdict['name'] in [value]:
302                 return rdict
303               
304         return tempdict
305
306
307     def filter(self, tagname, attribute, blacklist = [], whitelist = [], dom = None):
308         """
309         Removes all elements where:
310         1. tagname matches the element tag
311         2. attribute matches the element attribte
312         3. attribute value is in valuelist  
313         """
314
315         tempdict = {}
316         if not dom:
317             dom = self.rootNode
318        
319         if dom.localName in [tagname] and dom.attributes.has_key(attribute):
320             if whitelist and dom.attributes.get(attribute).value not in whitelist:
321                 dom.parentNode.removeChild(dom)
322             if blacklist and dom.attributes.get(attribute).value in blacklist:
323                 dom.parentNode.removeChild(dom)
324            
325         if dom.hasChildNodes():
326             for child in dom.childNodes:
327                 self.filter(tagname, attribute, blacklist, whitelist, child) 
328
329
330     def validateDicts(self):
331         types = {
332             'EInt' : int,
333             'EString' : str,
334             'EByteArray' : list,
335             'EBoolean' : bool,
336             'EFloat' : float,
337             'EDate' : date}
338
339
340     def pprint(self, r = None, depth = 0):
341         """
342         Pretty print the dict
343         """
344         line = ""
345         if r == None: r = self.dict
346         # Set the dept
347         for tab in range(0,depth): line += "    "
348         # check if it's nested
349         if type(r) == dict:
350             for i in r.keys():
351                 print line + "%s:" % i
352                 self.pprint(r[i], depth + 1)
353         elif type(r) in (tuple, list):
354             for j in r: self.pprint(j, depth + 1)
355         # not nested so just print.
356         else:
357             print line + "%s" %  r
358     
359
360
361 class RecordSpec(Rspec):
362
363     root_tag = 'record'
364     def parseDict(self, rdict, include_doc = False):
365         """
366         Convert a dictionary into a dom object and store it.
367         """
368         self.rootNode = self.dict2dom(rdict, include_doc)
369
370     def dict2dom(self, rdict, include_doc = False):
371         record_dict = rdict
372         if not len(rdict.keys()) == 1:
373             record_dict = {self.root_tag : rdict}
374         return Rspec.dict2dom(self, record_dict, include_doc)
375
376         
377 # vim:ts=4:expandtab
378