add remove_element() method
[sfa.git] / sfa / rspecs / rspec.py
1 #!/usr/bin/python 
2 from lxml import etree
3 from StringIO import StringIO
4 from datetime import datetime, timedelta
5 from sfa.util.xrn import *
6 from sfa.util.plxrn import hostname_to_urn
7 from sfa.util.config import Config  
8 from sfa.util.faults import SfaNotImplemented, InvalidRSpec
9
10 class RSpec:
11     header = '<?xml version="1.0"?>\n'
12     template = """<RSpec></RSpec>"""
13     config = Config()
14     xml = None
15     type = None
16     version = None
17     namespaces = None
18     user_options = {}
19   
20     def __init__(self, rspec="", namespaces={}, type=None, user_options={}):
21         self.type = type
22         self.user_options = user_options
23         if rspec:
24             self.parse_rspec(rspec, namespaces)
25         else:
26             self.create()
27
28     def create(self):
29         """
30         Create root element
31         """
32         # eg. 2011-03-23T19:53:28Z 
33         date_format = '%Y-%m-%dT%H:%M:%SZ'
34         now = datetime.utcnow()
35         generated_ts = now.strftime(date_format)
36         expires_ts = (now + timedelta(hours=1)).strftime(date_format) 
37         self.parse_rspec(self.template, self.namespaces)
38         self.xml.set('expires', expires_ts)
39         self.xml.set('generated', generated_ts)
40     
41     def parse_rspec(self, rspec, namespaces={}):
42         """
43         parse rspec into etree
44         """
45         parser = etree.XMLParser(remove_blank_text=True)
46         try:
47             tree = etree.parse(rspec, parser)
48         except IOError:
49             # 'rspec' file doesnt exist. 'rspec' is proably an xml string
50             try:
51                 tree = etree.parse(StringIO(rspec), parser)
52             except:
53                 raise InvalidRSpec('Must specify a xml file or xml string. Received: ' + rspec )
54
55         self.xml = tree.getroot()  
56         if namespaces:
57            self.namespaces = namespaces
58
59     def add_attribute(self, elem, name, value):
60         """
61         Add attribute to specified etree element    
62         """
63         opt = etree.SubElement(elem, name)
64         opt.text = value
65
66     def add_element(self, name, attrs={}, parent=None, text=""):
67         """
68         Generic wrapper around etree.SubElement(). Adds an element to 
69         specified parent node. Adds element to root node is parent is 
70         not specified. 
71         """
72         if parent == None:
73             parent = self.xml
74         element = etree.SubElement(parent, name)
75         if text:
76             element.text = text
77         if isinstance(attrs, dict):
78             for attr in attrs:
79                 element.set(attr, attrs[attr])  
80         return element
81
82     def remove_attribute(self, elem, name, value):
83         """
84         Removes an attribute from an element
85         """
86         if elem is not None:
87             opts = elem.iterfind(name)
88             if opts is not None:
89                 for opt in opts:
90                     if opt.text == value:
91                         elem.remove(opt)
92
93     def remove_element(self, element_name, root_node = None):
94         """
95         Removes all occurences of an element from the tree. Start at 
96         specified root_node if specified, otherwise start at tree's root.   
97         """
98         if not root_node:
99             root_node = self.xml
100
101         elements = root_node.xpath('//rspecv2:%s | //%s' % (element_name, element_name), \
102                                                             self.namespaces)
103         for element in elements:
104             parent = element.getparent()
105             parent.remove(element)
106          
107
108     def merge(self, in_rspec):
109         pass
110
111     def validate(self, schema):
112         """
113         Validate against rng schema
114         """
115         
116         relaxng_doc = etree.parse(schema)
117         relaxng = etree.RelaxNG(relaxng_doc)
118         if not relaxng(self.xml):
119             error = relaxng.error_log.last_error
120             message = "%s (line %s)" % (error.message, error.line)
121             raise InvalidRSpec(message)
122         return True
123         
124
125     def __str__(self):
126         return self.toxml()
127
128     def toxml(self):
129         return self.header + etree.tostring(self.xml, pretty_print=True)  
130         
131     def save(self, filename):
132         f = open(filename, 'w')
133         f.write(self.toxml())
134         f.close()
135  
136 if __name__ == '__main__':
137     rspec = RSpec()
138     print rspec