source: TI05-delivery/ows_framework/trunk/ows_server/ows_server/models/xmlHandler2.py @ 2834

try:
    from xml.etree import cElementTree as ET
    from xml.etree import ElementTree as pET
except ImportError:
    try:
        import cElementTree as ET
        import ElementTree as pET
    except ImportError:
        # For some reason when I install ElementTree with easyinstall it
        # is called "elementree".
        import elementtree.ElementTree as ET
        pET=ET
        
from xml.parsers.expat import ExpatError
import StringIO, re
XMLHDR='<?xml version="1.0"'

class xmlHandler:

    def __init__(self,xml,string=0):
        ''' Open an xml file (or string) and
           - if necessary correct nasty characters and/or orphans before passing to ET
           - load up an element-tree
           - collect a namespace map '''
        
        self.r1=None   # we only use the regex if we need them

        if string:
            self.xmls=xml
            xmlf=None
        else:
            xmlf=xml
            self.xmls=file(xmlf,'r').read()
        
        # Unfortunately we never know whether the incoming xml content is
        # unicode or a string ...
        try:
            self.xmls=unicode(self.xmls,'utf-8')
        except:
            pass
        self.xmls=self.xmls.encode('utf-8','replace')
        
        self.realns={}            
        self.__getns()
        
        try:
            self.tree=ET.XML(self.xmls)
        except SyntaxError:
            self.xmls=self.__fixXML(self.xmls)
            self.tree=ET.XML(self.xmls)
        
            
        self.__updatens()
    
    def __getns(self):
        
        ''' Get what the user intended out of elementtree namespaces '''
        #ought to do this with a regular expression, but needs must
        # or bettter yet, use iterparse in the first place, but that seemed slow.
        

        if self.xmls[0:19]==XMLHDR:
            self.root=1
            hb1=self.xmls.find('>')+1
        else:
            hb1=0
            self.root=0
        
        hb=self.xmls[hb1:].find('<')+1,self.xmls[hb1:].find('>')
        s=self.xmls[hb1+hb[0]:hb1+hb[1]]
        for w in s.split():
            n=w.split('=')
            if n[0]=='xmlns':
                self.realns[n[1][1:-1]]='default'
                self.defns=n[1][1:-1]#self.realns[n[1][1:-1]]
            elif n[0][0:6]=='xmlns:':
                self.realns[n[1][1:-1]]=n[0][6:]
            
    def tohtml(self):
        '''Lightweight HTML pretty printing of elementTree elements 
           and formatted using a css something like this:
            ===
            DIV.xmlElem {PADDING-LEFT: 20px;}
            .xmlAttrVal {COLOR:Red; }
            .xmlAttrTyp {COLOR:Green; }
            .xmlElemTag {COLOR:Blue; }
        .   highlight {BACKGROUND-COLOR:Yellow; }
            ===
            Line number is not yet implemented.
            '''
        lt,gt='<b>&lt;</b>','<b>&gt;</b>'
        def span(x,c): return '<span class="%s">%s</span>'%(c,x)
        def div(x,c): return '<div class="%s">%s</div>'%(c,x)
        def fix(x): 
            if x is None: return ''
            return x
        def et2html(elem):    
            strAttrib=''
            for att in elem.attrib:
                strAttrib+=' %s="%s"'%(span(att,'xmlAttrTyp'),span(elem.attrib[att],'xmlAttrVal'))
            result='%s%s%s%s%s'%(lt,span(elem.tag,"xmlElemTag"),strAttrib,gt,fix(elem.text))
            children=len(elem)
            if children:
                for item in elem:
                    result+=et2html(item)
                result+='%s%s/%s%s'%(fix(elem.tail),lt,span(elem.tag,'xmlElemTag'),gt)
            else:
                result+='%s/%s%s'%(lt,span(elem.tag,'xmlElemTag'),gt)
            return div(result,'xmlElem')
            
        ss=et2html(self.tree)
        h=''
        if self.root:h='%s%s %s="%s" %s="%s"%s'%(
            lt,'?xml',span('version','xmlAttrTyp'),'1.0',span('encoding','xmlAttrTyp'),'utf-8',gt)
        ss=self.__fixXML(ss)
        if self.realns=={}: 
            r=h+ss
        else: r= h+self.__nsfixpretty(ss,span)
        return '<div class="xmlDoc">%s</div>'%r
          
    def __nsfixpretty(self,s,span):
        ''' Yet another careful fix '''
        for ns in self.realns:
            r='{%s}'%ns
            if self.realns[ns]=='default':
                s=s.replace(r,'')
            else:
                s=s.replace(r,'%s:'%self.realns[ns])
        if self.realns=={}: return s
        # at this point we have no namespace list at the top
        rightArrow=s.find('</span>') # this is just after the tag, where we do want the namespace list
        nslist={} 
        for ns in self.realns: nslist[self.realns[ns]]=ns
        r=' %s="%s"'%(span('xmlns','xmlAttrTyp'),span(nslist['default'],'xmlAttrVal'))
        for ns in nslist:
            if ns<>'default': r+=' %s="%s"'%(span('xmlns:%s'%ns,'xmlAttrTyp'),span(nslist[ns],'xmlAttrVal'))
        return s[:rightArrow]+r+s[rightArrow:]
        
    def __updatens(self):
        ''' Update the element tree namespace map with our own map '''
        # *c*ElementTree doesn't have this update method (or at
        # least I can't find it), so you have to import ElementTree and call it on
        # that, then it all mysteriously works in cElementTree...

        pET._namespace_map.update(self.realns)

          
    def __str__(self):
        ### actually we should consider whether this was in the input or not 
        h=''
        if self.root:h='<?xml version="1.0" encoding="utf-8">'
        ss=ET.tostring(self.tree)
        ### ugly as sin, what happens if default: is in the text? We really ought to do this
        #properly in iterparse on loading the thing ...
        if self.realns=={}: return h+ss
        return self.__fixns(h,ss)
        
    def __fixns(self,h,ss):
        ''' Fix the namespaces after ET has produced a string '''
        ss=ss.replace('default:','')
        for ns in self.realns:
            r='xmlns:%s="%s"'%(self.realns[ns],ns)
            ss=ss.replace(r,'')
        #now fix the namespaces back in the first element
        rightArrow=ss.find('>')
        #reorder dictionary (I know I didn't need to do it but
        #code readability is worth a millisecond or two.
        nslist={}
        for ns in self.realns: nslist[self.realns[ns]]=ns
        r='xmlns="%s"'%nslist['default']
        for ns in nslist:
            if ns<>'default': r+=' xmlns:%s="%s"'%(ns,nslist[ns])
        h+=ss[:rightArrow]+r+ss[rightArrow:]
        return h
     
    def __fixXML(self,s):
        #first those nasty ampersands
        s=re.sub(r'&(?!\w+;)', '&amp;', s)
        #and now orphan > < signs
        if self.r1 is None:
            self.r1=re.compile('<([^>]*(<|$))')
            self.r2=re.compile('((^|>)[^<]*)>')
        old=''
        while s != old:
            old=s
            s=self.r1.sub(r'&lt;\1',s)
            s=self.r2.sub(r'\1&gt;',s)
        return s
        
    def _distributens(self,xpathExpression):
        ''' Actually we only support tag finding in this '''
        tags=xpathExpression.split('/')
        new=''
        for t in tags: 
            if t[1]<>'{': 
                new+='{%s}%s/'%(self.defns,t)
            else:
                new+=t+'/'
        new=new[0:-1]
        return new
        
    def getText(self,xpathExpression,multiple=0):
        ''' Get a text object sensibly, given ET API for xml doesn't handle 
        namespaces gracefully '''
        elem=self.tree 
        if multiple:
                r=elem.findall(self._distributens(xpathExpression))
        else:
                r=[elem.find(self._distributens(xpathExpression)),]
        try:  # if element is None, this should fail ...
                rr=[]
                for i in r:
                    t=i.text 
                    if t is not None: 
                        rr.append(t)
                    else: rr.append('')
        except:
                rr=['',]
        if multiple: 
                return rr
        else: return rr[0]
          
          
if __name__=="__main__":
    
    import unittest
    
    class TestCase(unittest.TestCase):  
        
        def setup(self):
            self.ss='''<?xml version="1.0" encoding="UTF-8"?>
                <Dataset xmlns:swe="http://www.opengis.net/swe" xmlns:gml="http://www.opengis.net/gml"
                 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns:moles="http://ndg.nerc.ac.uk/moles"
                 xmlns:om="http://www.opengis.net/om" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns="http://ndg.nerc.ac.uk/csml" 
                 id="FGPfF9i0"><CSMLFeatureCollection gml:id="AfEj15o6"/><om:blah>blahvalue</om:blah><foo>foovalue</foo></Dataset>'''
                 
        def testns(self):
            ''' Make sure we extract the namespaces correctly '''
            self.setup()
            x=xmlHandler(self.ss,string=1)
            self.assertEqual(x.realns,{'http://www.opengis.net/om':'om', 'http://www.opengis.net/gml':'gml', 
            'http://ndg.nerc.ac.uk/csml':'default', 'http://www.opengis.net/swe':'swe', 
            'http://www.w3.org/1999/xlink':'xlink', 
            'http://www.w3.org/2001/XMLSchema-instance':'xsi', 'http://ndg.nerc.ac.uk/moles':'moles'})
            
        def teststr(self):
            ''' Make sure we can get a string version after loading '''
            self.setup()
            x=xmlHandler(self.ss,string=1)
            self.assertEqual('<?xml version="1.0" encoding="utf-8"><Dataset id="FGPfF9i0" xmlns="http://ndg.nerc.ac.uk/csml" xmlns:om="http://www.opengis.net/om" xmlns:gml="http://www.opengis.net/gml" xmlns:swe="http://www.opengis.net/swe" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns:moles="http://ndg.nerc.ac.uk/moles"><CSMLFeatureCollection gml:id="AfEj15o6"  /><om:blah >blahvalue</om:blah><foo>foovalue</foo></Dataset>',str(x))
            
        def testorphans(self):
            ''' Make sure we can handle orphan characters properly '''
            s='<data> 1<2</data>'
            x=xmlHandler(s,string=1)
            self.assertEqual('<data> 1&lt;2</data>',str(x))

        def testAmpsersand1(self):
            ''' Can we load unescaped ampersands?'''
            s='<data> a & b </data>'
            x=xmlHandler(s,string=1)
            self.assertEqual('<data> a &amp; b </data>',str(x))
            
        def testAmpersand2(self):
            ''' Do we output proper things? '''
            s='<data> 2 &amp; 3 &lt; 8 </data>'
            x=xmlHandler(s,string=1)
            self.assertEqual('<data> 2 &amp; 3 &lt; 8 </data>',str(x))
            
        def testPrettyPrint(self):
            ''' Test a simple pretty print '''
            s='<?xml version="1.0" encoding="utf-8"?><data><element>stuff</element></data>'
            x=xmlHandler(s,string=1)
            h=x.tohtml() # only testing the mechanics, not the result
            
            
        #turn off the test
        def AtestRealDIF(self):
            ''' Test a real DIF from the ndgRetrieve stable '''
            f='ndgRetrieve.badc.nerc.ac.uk__DIF__dataent_11738019833217179.debug.xml'
            x=xmlHandler(f)
            y=str(x)  # only testing the mechanics, not the result
            
            
        def testDIF(self):
            s='''<DIF xmlns="http://gcmd.gsfc.nasa.gov/Aboutus/xml/dif/" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"><Entry_ID>badc.nerc.ac.uk__DIF__dataent_11738019833217179</Entry_ID></DIF>'''
            x=xmlHandler(s,string=1)
            print x.realns
            print str(x)
            h=x.tohtml()
            print h
            
    unittest.main()