File contents

"""	Class to handle VOTable
	Created: 2005-05-31 by Shui Hung Kwok, shkwok at computer.org
	Edited: 2007-02-12 by Noel Winstanley, noel.winstanley at manchester.ac.uk
		- added ability to parse strings as votable-xml
	See http://www.ivoa.net/Documents/latest/VOT.html .
	
"""

import sys
from types import *
import xml.sax
import xml.sax.handler 
import StringIO

class VONode (object):
	"""	Class representing an XML node of a VOTable
	""" 
	def __init__ (self, tagname=('','')):
		self._tagname = tagname
		self._nodeList = []
		self._attributes = {}

	def addNode (self, node):
		self._nodeList.append (node)
		if not isinstance (node, (StringType, UnicodeType)):
			name = node.getNamePart ()
			try:
				val = self.__dict__[name]
				if isinstance (val, ListType):
					val.append (node)
				else:
					self.__dict__[name] = [val, node]
			except:
				self.__dict__[name] = node
		else:
			self.content = node

	def addAttribute (self, attr):
		name,value = attr
		self._attributes[name] = value
		self.__dict__[name] = value

	def addAttributes (self, attrs):
		self._attributes.update (attrs)
		for k,v in attrs.items ():
			self.__dict__[k[1]] = v

	def __getItem__ (self, idx):
		return self._nodeList[idx].getContent ()

	def getAttribute (self, name):
		""" Returns attribute by name or '' if not found """
		return self._attributes.get (name)

	def getAttributes (self):
		""" Returns all attributes.
		"""
		res = {}
		for (ns,n),at in self._attributes.items():
			res[n] = at
		return res

	def getNodeList (self):
		""" Returns a list of nodes that are of type VONode
		"""
		res = []
		for node in self._nodeList:
			try:
				l = node._nodeList
				res.append (node)
			except Exception, e:
				pass
		return res		

	def getContent (self):
		""" Returns all strings of the node.
		"""
		res = []
		for node in self._nodeList:
			try:
				l = node.lower ()
				res.append (node)
			except Exception, e:
				pass
		return ''.join (res)
	
	def getNamePart (self):
		try:
			ns, n = self._tagname
			return n
		except:
			return n
			
	def getNodesByName (self, look4):
		""" Returns a list of nodes whose tagname = look4
		"""
		res = []
		for node in self._nodeList:
			try:
				if look4 != node.getNamePart ():
					continue
				l = node._nodeList
				res.append (node)
			except Exception, e:
				pass
		return res		

	def __str__ (self):
		try:
			return self.content
		except:
			return self.buildName (self._tagname)

	def getNode (self, path):
		"""	Returns a node for a given path.
			Path is of the form /tag1/tag2/tag3.
			Path can include array index, like /tag1/tag2[3]/tag4.
		"""
		node = self
		children = []
		pathArray = path.split ("/")

		rootName = self.getNamePart ()
		if rootName != pathArray[1]:
			return None

		pathArray = pathArray[2:]
		for elem in pathArray:
			tmp = elem.replace ('[', ']')
			list = tmp.split (']')
			name = list[0]
			if len (list) > 1:
				idx = int (list[1])
			else:
				idx = 0
			children = node.getNodesByName (name)
			nr = len (children)
			if idx >= nr:
				return None
			node = children [idx]
		return node	

	def getNodesByPath (self, path):
		"""	Returns an array of VONodes for a given path.
			Path is of the form /tag1/tag2/tag3.
			Path can include array index, like /tag1/tag2[3]/tag4.
		"""
		node = self
		children = []
		pathArray = path.split ("/")

		rootName = self.getNamePart ()
		if rootName != pathArray[1]:
			return None

		pathArray = pathArray[2:]
		for elem in pathArray:
			tmp = elem.replace ('[', ']')
			list = tmp.split (']')
			name = list[0]
			if len (list) > 1:
				idx = int (list[1])
			else:
				idx = 0
			children = node.getNodesByName (name)
			nr = len (children)
			if idx >= nr:
				return None
			node = children [idx]
		return children	

	def buildName (self, tname):
		""" Returns a name with namespace as prefix
			or just name if no namespace
			Note that the prefix is the real namespace
			and not the abbreviation used in the original XML
		"""
		ns,n = tname
		"""
		if ns:
			return "%s:%s" % (self.qname, n)
		else: 
			return n
		"""
		return n

	def printAllNodes (self, func=sys.stdout.write, prefix=''):
		"""	Recursive method to visit all nodes of the tree
			and calls the provided function to output the content.
		"""
		func ("%s<%s" % (prefix, self.buildName (self._tagname)))
		for ns,v in self._attributes.items ():
			func (" %s='%s'" % (self.buildName ((ns)), v))
		func (">")
		
		last = 0
		for n in self._nodeList:
			if isinstance (n, (StringType, UnicodeType)):
				if last == 2:
					func ("\n%s" % prefix)
				func ("%s" % n)
				last = 1
			else:
				if last <= 1:
					func ("\n")
				n.printAllNodes (func, prefix + '   ')
				last = 2
		if last <= 1:
			func ("</%s>\n" % self.buildName (self._tagname))
		else:
			func ("%s</%s>\n" % (prefix, self.buildName (self._tagname)))
		
class VOXMLHandler (xml.sax.handler.ContentHandler):
	""" Class implementing callbacks for the SAX parser.
	"""
	def __init__ (self, vonode=VONode):
		# Create a parser
		xml.sax.handler.ContentHandler.__init__ (self)
		self.parser = xml.sax.make_parser()
		self.parser.setFeature (xml.sax.handler.feature_namespaces, 1)
		self.parser.setContentHandler (self)
		self.vonode = vonode
		self.sentinel = vonode ()
		self.currNode = self.sentinel
		self.stack = []

	def startElementNS (self, (urn, name), qname, attrs):
		#print "start ", name
		self.stack.append (self.currNode)

		self.currNode = self.vonode ((urn, name))
		self.currNode.addAttributes (attrs)

	def characters (self, chunk):
		buf = chunk.strip ()
		if len (buf) == 0: return
		self.currNode.addNode (buf)

	def endElementNS (self, (urn, name), qname):
		#print "end ", name
		newNode = self.currNode
		self.currNode = self.stack.pop ()
		self.currNode.addNode (newNode)

	def parse (self, source):
		""" Main entry point.
			Source can be URL or file name.
		"""
		self.parser.parse (source)
		return self.sentinel 
	#NWW
	def parseText(self,t):
		""" Alternate entry point. Text is some xml to be parsed"""
		inputSource = xml.sax.xmlreader.InputSource()
		inputSource.setByteStream(StringIO.StringIO(t))
		self.parser.parse(inputSource)
		return self.sentinel
	
class VOXML (object):
	"""	Class to implement XML parser
	"""
	def __init__ (self, source=None, vonode=VONode):
		""" Instantiate a VOXML
			source can be URL, file name or a string representing the VOXML.
			vonode is a class representing VONode, must be derived from or
			compatible with VONode.
		"""
		self.vonode = vonode 
		self.root = None
		if source != None:
			self.parse (source)

	def parse (self, source):
		""" Invokes XML parser and stores VOXML
			in self.root as VONode.
		"""
		parser = VOXMLHandler (self.vonode)
		self.root = parser.parse (source)
	#NWW
	def parseText(self,text):
		""" Invokes XML parser on inlined xml text"""
		parser = VOXMLHandler(self.vonode)
		self.root = parser.parseText(text)

	def printAllNodes (self, func=sys.stdout.write):	
		"""	Output entire content as XML.
			func is the output method, defined as:
				func (outString)
		"""
		node = self.root._nodeList[0] 
		node.printAllNodes (func)

	def getNode (self, path):
		""" Returns a VONode of the given path.
		"""
		return self.root._nodeList[0].getNode (path)

	def getContent (self, path):
		""" Returns the content of a node.
			Only strings are returned.
		"""
		node = self.getNode (path)
		return node.getContent ()

class VOTable (VOXML):
	""" Implementation of VOTable 
	"""

	def __init__ (self, source=None, vonode=VONode):
		""" Contruct an instance of VOTable
		"""
		VOXML.__init__ (self, source, vonode)

	def getColumnIdx (self, val):
		""" Returns the column index for the given name
			Will return any attribute value matching val.
		"""
		fields = self.getFields ()
		for coln, f in enumerate (fields):
			if val in f._attributes.values ():
				return coln
		return -1
		
	def getFields (self):
		""" Returns a list of VONode representing all the fields
		"""
		#table = self.root.VOTABLE.RESOURCE.TABLE
		#return table.getNodesByName ('FIELD')
		return self.root.VOTABLE.RESOURCE.TABLE.FIELD
	
	def getParams (self):
		"""	Returns a list of VONode representing all PARAMS
		"""
		return self.root.VOTABLE.RESOURCE.RESOURCE.PARAM

	def getFieldsAttrs (self):
		"""	Returns a list of maps that contains attributes.
			Returned list looks like this: [{},{},...]
		"""
		res = []
		fields = self.getFields ()
		for elem in fields:
			try:
				res.append (elem.getAttributes())
			except:
				pass
		return res
	
	def getDataRows (self):
		"""	Returns a list of VONodes representing rows of the table.
			Use getData () to extract data from each row.
			for x in getDataRows ():
				data = getData (x)
				#data = [values ...]
		"""
		tableData = self.root.VOTABLE.RESOURCE.TABLE.DATA.TABLEDATA
		return tableData._nodeList

	def getData (self, row):
		""" row is a VONode <TR> parent of a list of <TD>.
			Returns a list of values.
		"""
		res = []
		list = row._nodeList
		for elm in list:
			try:
				res.append (elm.getContent ())
			except:
				res.append ('')
		return res	
	
	def append (self, vot):
		""" Appends votable vot to the end of this VOTable. 
			No tests to see if fields are the same. 
			vot must have the same fields. 
		"""
		try:
			node1 = self.root.VOTABLE.RESOURCE.TABLE.DATA.TABLEDATA
		except:
			node1 = None
		try:
			node2 = vot.root.VOTABLE.RESOURCE.TABLE.DATA.TABLEDATA
		except:
			node2 = None

		if node1:
			if node2:
				node1._nodeList.extend (node2._nodeList)
		else:
			if node2:
				self.root.VOTABLE.RESOURCE.TABLE.DATA.TABLEDATA = node2
		
	def output2CSV (self, write=sys.stdout.write):
		""" Outputs the content in CSV format
			with fields names as column heads.
		"""
		fields = self.getFieldsAttrs ()
		if not fields: return
		headers = []
		for colnr,fd in enumerate (fields):
			hdr = "col%d" % colnr
			for i in ['name', 'ID', 'id']:
				try:
					hdr = fd[i]
					break
				except:
					pass
			headers.append (hdr)
	
		write (','.join (headers))
	
		rows = self.getDataRows ()
		getData = self.getData # get a shortcut for the method
		for row in rows:
			write (','.join (getData (row)))

if __name__ == '__main__':
	votable = VOTable()
	votable.parse (sys.argv[1])
	votable.printAllNodes ()
	#print [x.getAttribute ('ID') for x in votable.getFields () ]
	#print votable.root.VOTABLE.RESOURCE.TABLE.DATA.TABLEDATA.TR[1].TD[1]
	#print votable.getFields ()