Le portail francophone de la géomatique

from qgis.core import *
from qgis.gui import *

import logging
logging.getLogger('fastkml.config').addHandler(logging.NullHandler())
from fastkml import kml

import shapely.wkt
import os
import re
import zipfile
from xml.etree import ElementTree

class KmlExport(object):
    def __init__(self):
        self.ns='{http://www.opengis.net/kml/2.2}'
        ElementTree.register_namespace('',"http://www.opengis.net/kml/2.2")
        
            
    def projectToKml(self, projectId, outputDirectory, prettyprint=False):
        
        root=QgsProject.instance().layerTreeRoot()
        k = kml.KML()
        d = kml.Document(self.ns, projectId, projectId)
        k.append(d)
        
        self._addLayersInFolder(root, d, "")
        
        self._outputToKmlAndKmz(outputDirectory, k, projectId, prettyprint)

    def _addLayersInFolder(self,folder,kmlContainer, baseName):
        for node in folder.children():
            if QgsLayerTree.isGroup(node):
                folderName=node.name()
                folderName=baseName+folderName
                f = kml.Folder(self.ns, folderName, folderName)
                kmlContainer.append(f)
                self._addLayersInFolder(node, f, folderName+"_")
            else:
                layerName=node.layerName()
                layerName=baseName+layerName
                f = kml.Folder(self.ns, layerName, layerName)
                kmlContainer.append(f)
                self._addFeaturesInLayer(node.layer(), f)
                
    def _addFeaturesInLayer(self, layer, kmlContainer):
        nameField=self._nameField(layer.name())
        idLayer="#%s"%self._toId(layer.name())

        fieldNames=map(lambda x:x.name(),layer.dataProvider().fields().toList())
        nameFieldIndex=layer.fieldNameIndex(nameField)
        for feature in layer.getFeatures():
            try:
                attributes=feature.attributes()
                geometry=feature.geometry()
                if not geometry:
                    qDebug("Invalid geometry in layer %s (id='%s',geometry='%s')"%(layer.name(),
                                            attributes[nameFieldIndex],geometry.exportToWkt()))
                    continue
                
                p=kml.Placemark(self.ns)
                p.name=attributes[nameFieldIndex]
                p.geometry=shapely.wkt.loads(geometry.exportToWkt())
                data=self._dataElementsFromAttributes(fieldNames,attributes )
                if len(data) > 0:
                    p.extended_data = kml.ExtendedData(self.ns, [kml.SchemaData(self.ns, idLayer, data)])
                kmlContainer.append(p)
            except Exception:
                qDebug("Error in layer %s (id='%s',geometry='%s')"%(layer.name(),
                                                                    attributes[nameFieldIndex],geometry.exportToWkt()))
                raise
            
            
    def _outputToKmlAndKmz(self, outputDirectory, kmlData, base, prettyprint):
        if prettyprint:
            # etree does not support prettyprint. Do it ourselves
            content=self._prettyPrint(kmlData.etree_element())
        else:
            content=kmlData.to_string().encode('utf-8')
        kmlFileName=base+".kml"
        with open(os.path.join(outputDirectory,kmlFileName),"wb") as f:
            f.write(content)
            
        kmzFileName=base+".kmz"
        with zipfile.ZipFile(os.path.join(outputDirectory,kmzFileName), 'w', zipfile.ZIP_DEFLATED) as kmz:
            kmz.writestr("content.kml", content)
        
    def _nameField(layerName):
        # FIXME add some logic
        return "ID"
        
    def _dataElementsFromAttributes(self, fieldNames, attributes, removeNullValues=False):
        data=zip(fieldNames,attributes)
        if removeNullValues:
            data=filter(lambda x:x[1] <> None, data)
        return map(lambda x:(x[0], self._toKmlValue(x[1])),data)
    
    def _toKmlValue(self, value):
        if isinstance(value, str) or isinstance(value, unicode):
            return value
        else:
            return str(value)

    def _toId(self, name):
        return re.sub(r'[^a-zA-Z_0-9\-]','_',name)
    
    def _prettyPrint(self, elem):
        self._indent(elem)
        return ElementTree.tostring(elem,"utf-8")

    def _indent(self,elem, level=0):
        i = "\n" + level*"  "
        if len(elem):
            if not elem.text or not elem.text.strip():
                elem.text = i + "  "
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
            for elem in elem:
                self._indent(elem, level+1)
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
        else:
            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = i