Added file output to gdal extractprojection, added exception for non-gdal sources and fixed boolean parameter handling + some formatting

havatv · havatv · commit 8976f6822f1e · 2018-11-01T16:03:43.000+01:00
diff --git a/python/plugins/processing/algs/gdal/extractprojection.py b/python/plugins/processing/algs/gdal/extractprojection.py
@@ -32,33 +32,46 @@
 from osgeo import gdal, osr
 
 from processing.algs.gdal.GdalAlgorithm import GdalAlgorithm
+from qgis.core import QgsProcessingException
 from qgis.core import (QgsProcessingParameterRasterLayer,
-                       QgsProcessingParameterBoolean)
+                       QgsProcessingParameterBoolean,
+                       QgsProcessingOutputFile)
 
 pluginPath = os.path.split(os.path.split(os.path.dirname(__file__))[0])[0]
 
 
 class ExtractProjection(GdalAlgorithm):
 
     INPUT = 'INPUT'
+    PRJ_FILE_CREATE = 'PRJ_FILE_CREATE'
+    WORLD_FILE = 'WORLD_FILE'
     PRJ_FILE = 'PRJ_FILE'
 
     def __init__(self):
         super().__init__()
 
     def initAlgorithm(self, config=None):
-        self.addParameter(QgsProcessingParameterRasterLayer(self.INPUT, self.tr('Input file')))
-        self.addParameter(QgsProcessingParameterBoolean(self.PRJ_FILE,
-                                           self.tr('Create also .prj file'), False))
-
+        self.addParameter(QgsProcessingParameterRasterLayer(
+                              self.INPUT,
+                              self.tr('Input file')))
+        self.addParameter(QgsProcessingParameterBoolean(
+                              self.PRJ_FILE_CREATE,
+                              self.tr('Create also .prj file'), False))
+        self.addOutput(QgsProcessingOutputFile(self.WORLD_FILE,
+                                               self.tr('World file')))
+        self.addOutput(QgsProcessingOutputFile(
+                              self.PRJ_FILE,
+                              self.tr('ESRI Shapefile prj file')))
+        
     def name(self):
         return 'extractprojection'
 
     def displayName(self):
         return self.tr('Extract projection')
 
     def icon(self):
-        return QIcon(os.path.join(pluginPath, 'images', 'gdaltools', 'projection-export.png'))
+        return QIcon(os.path.join(pluginPath, 'images', 'gdaltools',
+                     'projection-export.png'))
 
     def group(self):
         return self.tr('Raster projections')
@@ -69,12 +82,19 @@ def groupId(self):
     def commandName(self):
         return 'extractprojection'
 
-    def getConsoleCommands(self, parameters, context, feedback, executing=True):
+    def getConsoleCommands(self, parameters, context, feedback,
+                           executing=True):
         return [self.commandName()]
 
     def processAlgorithm(self, parameters, context, feedback):
-        createPrj = QgsProcessingParameterBoolean(self.PRJ_FILE)
-        raster = self.parameterAsRasterLayer(parameters, self.INPUT, context)
+        createPrj = self.parameterAsBool(parameters,
+                                         self.PRJ_FILE_CREATE,
+                                         context)
+        raster = self.parameterAsRasterLayer(parameters, self.INPUT,
+                                             context)
+        if not raster.dataProvider().name() == 'gdal':
+            raise QgsProcessingException('This algorithm can only '\
+                                         'be used with GDAL raster layers')
         rasterPath = raster.source()
         rasterDS = gdal.Open(rasterPath, gdal.GA_ReadOnly)
         geotransform = rasterDS.GetGeoTransform()
@@ -85,6 +105,7 @@ def processAlgorithm(self, parameters, context, feedback):
 
         outFileName = os.path.splitext(str(rasterPath))[0]
 
+        results = {}
         if crs != '' and createPrj:
             tmp = osr.SpatialReference()
             tmp.ImportFromWkt(crs)
@@ -94,6 +115,9 @@ def processAlgorithm(self, parameters, context, feedback):
 
             with open(outFileName + '.prj', 'wt') as prj:
                 prj.write(crs)
+            results[self.PRJ_FILE] = outFileName + '.prj'
+        else:    
+            results[self.PRJ_FILE] = None
 
         with open(outFileName + '.wld', 'wt') as wld:
             wld.write('%0.8f\n' % geotransform[1])
@@ -106,5 +130,6 @@ def processAlgorithm(self, parameters, context, feedback):
             wld.write('%0.8f\n' % (geotransform[3] +
                                    0.5 * geotransform[4] +
                                    0.5 * geotransform[5]))
-        return {}
+        results[self.WORLD_FILE] = outFileName + '.wld'
 
+        return results
