sfa/util/xml.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.faults import SfaNotImplemented, InvalidXML
   8
   9 class XpathFilter:
  10     @staticmethod
  11     def xpath(filter={}):
  12         xpath = ""
  13         if filter:
  14             filter_list = []
  15             for (key, value) in filter.items():
  16                 if key == 'text':
  17                     key = 'text()'
  18                 else:
  19                     key = '@'+key
  20                 if isinstance(value, str):
  21                     filter_list.append('%s="%s"' % (key, value))
  22                 elif isinstance(value, list):
  23                     filter_list.append('contains("%s", %s)' % (' '.join(map(str, value)), key))
  24             if filter_list:
  25                 xpath = ' and '.join(filter_list)
  26                 xpath = '[' + xpath + ']'
  27         return xpath
  28
  29 class XML:
  30
  31     def __init__(self, xml=None):
  32         self.root = None
  33         self.namespaces = None
  34         self.default_namespace = None
  35         self.schema = None
  36         if isinstance(xml, basestring):
  37             self.parse_xml(xml)
  38         elif isinstance(xml, etree._ElementTree):
  39             self.root = xml.getroot()
  40         elif isinstance(xml, etree._Element):
  41             self.root = xml
  42
  43     def parse_xml(self, xml):
  44         """
  45         parse rspec into etree
  46         """
  47         parser = etree.XMLParser(remove_blank_text=True)
  48         try:
  49             tree = etree.parse(xml, parser)
  50         except IOError:
  51             # 'rspec' file doesnt exist. 'rspec' is proably an xml string
  52             try:
  53                 tree = etree.parse(StringIO(xml), parser)
  54             except Exception, e:
  55                 raise InvalidXML(str(e))
  56         self.root = tree.getroot()
  57         # set namespaces map
  58         self.namespaces = dict(self.root.nsmap)
  59         if 'default' not in self.namespaces and None in self.namespaces:
  60             self.namespaces['default'] = self.namespaces[None]
  61         # If the 'None' exist, then it's pointing to the default namespace. This makes
  62         # it hard for us to write xpath queries for the default naemspace because lxml
  63         # wont understand a None prefix. We will just associate the default namespeace
  64         # with a key named 'default'.
  65         if None in self.namespaces:
  66             default_namespace = self.namespaces.pop(None)
  67             self.namespaces['default'] = default_namespace
  68
  69         # set schema
  70         for key in self.root.attrib.keys():
  71             if key.endswith('schemaLocation'):
  72                 # schema location should be at the end of the list
  73                 schema_parts  = self.root.attrib[key].split(' ')
  74                 self.schema = schema_parts[1]
  75                 namespace, schema  = schema_parts[0], schema_parts[1]
  76                 break
  77
  78     def parse_dict(self, d, root_tag_name='xml', element = None):
  79         if element is None:
  80             self.parse_xml('<%s/>' % root_tag_name)
  81             element = self.root
  82
  83         if 'text' in d:
  84             text = d.pop('text')
  85             element.text = text
  86
  87         # handle repeating fields
  88         for (key, value) in d.items():
  89             if isinstance(value, list):
  90                 value = d.pop(key)
  91                 for val in value:
  92                     if isinstance(val, dict):
  93                         child_element = etree.SubElement(element, key)
  94                         self.parse_dict(val, key, child_element)
  95
  96         element.attrib.update(d)
  97
  98     def validate(self, schema):
  99         """
 100         Validate against rng schema
 101         """
 102         relaxng_doc = etree.parse(schema)
 103         relaxng = etree.RelaxNG(relaxng_doc)
 104         if not relaxng(self.root):
 105             error = relaxng.error_log.last_error
 106             message = "%s (line %s)" % (error.message, error.line)
 107             raise InvalidXML(message)
 108         return True
 109
 110     def xpath(self, xpath, namespaces=None):
 111         if not namespaces:
 112             namespaces = self.namespaces
 113         return self.root.xpath(xpath, namespaces=namespaces)
 114
 115     def set(self, key, value):
 116         return self.root.set(key, value)
 117
 118     def add_attribute(self, elem, name, value):
 119         """
 120         Add attribute to specified etree element
 121         """
 122         opt = etree.SubElement(elem, name)
 123         opt.text = value
 124
 125     def add_element(self, name, attrs={}, parent=None, text=""):
 126         """
 127         Generic wrapper around etree.SubElement(). Adds an element to
 128         specified parent node. Adds element to root node is parent is
 129         not specified.
 130         """
 131         if parent == None:
 132             parent = self.root
 133         element = etree.SubElement(parent, name)
 134         if text:
 135             element.text = text
 136         if isinstance(attrs, dict):
 137             for attr in attrs:
 138                 element.set(attr, attrs[attr])
 139         return element
 140
 141     def remove_attribute(self, elem, name, value):
 142         """
 143         Removes an attribute from an element
 144         """
 145         if elem is not None:
 146             opts = elem.iterfind(name)
 147             if opts is not None:
 148                 for opt in opts:
 149                     if opt.text == value:
 150                         elem.remove(opt)
 151
 152     def remove_element(self, element_name, root_node = None):
 153         """
 154         Removes all occurences of an element from the tree. Start at
 155         specified root_node if specified, otherwise start at tree's root.
 156         """
 157         if not root_node:
 158             root_node = self.root
 159
 160         if not element_name.startswith('//'):
 161             element_name = '//' + element_name
 162
 163         elements = root_node.xpath('%s ' % element_name, namespaces=self.namespaces)
 164         for element in elements:
 165             parent = element.getparent()
 166             parent.remove(element)
 167
 168     def attributes_list(self, elem):
 169         # convert a list of attribute tags into list of tuples
 170         # (tagnme, text_value)
 171         opts = []
 172         if elem is not None:
 173             for e in elem:
 174                 opts.append((e.tag, str(e.text).strip()))
 175         return opts
 176
 177     def get_element_attributes(self, elem=None, depth=0):
 178         if elem == None:
 179             elem = self.root_node
 180         if not hasattr(elem, 'attrib'):
 181             # this is probably not an element node with attribute. could be just and an
 182             # attribute, return it
 183             return elem
 184         attrs = dict(elem.attrib)
 185         attrs['text'] = str(elem.text).strip()
 186         attrs['parent'] = elem.getparent()
 187         if isinstance(depth, int) and depth > 0:
 188             for child_elem in list(elem):
 189                 key = str(child_elem.tag)
 190                 if key not in attrs:
 191                     attrs[key] = [self.get_element_attributes(child_elem, depth-1)]
 192                 else:
 193                     attrs[key].append(self.get_element_attributes(child_elem, depth-1))
 194         else:
 195             attrs['child_nodes'] = list(elem)
 196         return attrs
 197
 198     def merge(self, in_xml):
 199         pass
 200
 201     def __str__(self):
 202         return self.toxml()
 203
 204     def toxml(self):
 205         return etree.tostring(self.root, encoding='UTF-8', pretty_print=True)
 206
 207     def todict(self, elem=None):
 208         if elem is None:
 209             elem = self.root
 210         d = {}
 211         d.update(elem.attrib)
 212         d['text'] = elem.text
 213         for child in elem.iterchildren():
 214             if child.tag not in d:
 215                 d[child.tag] = []
 216             d[child.tag].append(self.todict(child))
 217         return d
 218
 219     def save(self, filename):
 220         f = open(filename, 'w')
 221         f.write(self.toxml())
 222         f.close()
 223
 224 if __name__ == '__main__':
 225     rspec = RSpec('/tmp/resources.rspec')
 226     print rspec
 227