Port export geometry info to new API

nyalldawson · nyalldawson · commit ab70e050a682 · 2017-07-27T11:21:24.000+10:00
diff --git a/python/plugins/processing/algs/qgis/ExportGeometryInfo.py b/python/plugins/processing/algs/qgis/ExportGeometryInfo.py
@@ -34,13 +34,12 @@
                        QgsField,
                        QgsWkbTypes,
                        QgsFeatureSink,
-                       QgsProcessingUtils,
-                       QgsDistanceArea)
+                       QgsDistanceArea,
+                       QgsProcessingParameterFeatureSource,
+                       QgsProcessingParameterEnum,
+                       QgsProcessingParameterFeatureSink)
 
 from processing.algs.qgis.QgisAlgorithm import QgisAlgorithm
-from processing.core.parameters import ParameterVector
-from processing.core.parameters import ParameterSelection
-from processing.core.outputs import OutputVector
 from processing.tools import vector
 
 pluginPath = os.path.split(os.path.split(os.path.dirname(__file__))[0])[0]
@@ -66,18 +65,16 @@ def __init__(self):
         self.export_z = False
         self.export_m = False
         self.distance_area = None
-
-    def initAlgorithm(self, config=None):
         self.calc_methods = [self.tr('Layer CRS'),
                              self.tr('Project CRS'),
                              self.tr('Ellipsoidal')]
 
-        self.addParameter(ParameterVector(self.INPUT,
-                                          self.tr('Input layer')))
-        self.addParameter(ParameterSelection(self.METHOD,
-                                             self.tr('Calculate using'), self.calc_methods, 0))
-
-        self.addOutput(OutputVector(self.OUTPUT, self.tr('Added geom info')))
+    def initAlgorithm(self, config=None):
+        self.addParameter(QgsProcessingParameterFeatureSource(self.INPUT,
+                                                              self.tr('Input layer')))
+        self.addParameter(QgsProcessingParameterEnum(self.METHOD,
+                                                     self.tr('Calculate using'), options=self.calc_methods, defaultValue=0))
+        self.addParameter(QgsProcessingParameterFeatureSink(self.OUTPUT, self.tr('Added geom info')))
 
     def name(self):
         return 'exportaddgeometrycolumns'
@@ -86,36 +83,36 @@ def displayName(self):
         return self.tr('Export/Add geometry columns')
 
     def processAlgorithm(self, parameters, context, feedback):
-        layer = QgsProcessingUtils.mapLayerFromString(self.getParameterValue(self.INPUT), context)
-        method = self.getParameterValue(self.METHOD)
+        source = self.parameterAsSource(parameters, self.INPUT, context)
+        method = self.parameterAsEnum(parameters, self.METHOD, context)
 
-        geometryType = layer.geometryType()
-        fields = layer.fields()
+        wkb_type = source.wkbType()
+        fields = source.fields()
 
-        if geometryType == QgsWkbTypes.PolygonGeometry:
+        if QgsWkbTypes.geometryType(wkb_type) == QgsWkbTypes.PolygonGeometry:
             areaName = vector.createUniqueFieldName('area', fields)
             fields.append(QgsField(areaName, QVariant.Double))
             perimeterName = vector.createUniqueFieldName('perimeter', fields)
             fields.append(QgsField(perimeterName, QVariant.Double))
-        elif geometryType == QgsWkbTypes.LineGeometry:
+        elif QgsWkbTypes.geometryType(wkb_type) == QgsWkbTypes.LineGeometry:
             lengthName = vector.createUniqueFieldName('length', fields)
             fields.append(QgsField(lengthName, QVariant.Double))
         else:
             xName = vector.createUniqueFieldName('xcoord', fields)
             fields.append(QgsField(xName, QVariant.Double))
             yName = vector.createUniqueFieldName('ycoord', fields)
             fields.append(QgsField(yName, QVariant.Double))
-            if QgsWkbTypes.hasZ(layer.wkbType()):
+            if QgsWkbTypes.hasZ(source.wkbType()):
                 self.export_z = True
                 zName = vector.createUniqueFieldName('zcoord', fields)
                 fields.append(QgsField(zName, QVariant.Double))
-            if QgsWkbTypes.hasM(layer.wkbType()):
+            if QgsWkbTypes.hasM(source.wkbType()):
                 self.export_m = True
                 zName = vector.createUniqueFieldName('mvalue', fields)
                 fields.append(QgsField(zName, QVariant.Double))
 
-        writer = self.getOutputFromName(self.OUTPUT).getVectorWriter(fields, layer.wkbType(), layer.crs(),
-                                                                     context)
+        (sink, dest_id) = self.parameterAsSink(parameters, self.OUTPUT, context,
+                                               fields, wkb_type, source.sourceCrs())
 
         coordTransform = None
 
@@ -126,14 +123,17 @@ def processAlgorithm(self, parameters, context, feedback):
 
         self.distance_area = QgsDistanceArea()
         if method == 2:
-            self.distance_area.setSourceCrs(layer.crs())
+            self.distance_area.setSourceCrs(source.sourceCrs())
             self.distance_area.setEllipsoid(context.project().ellipsoid())
         elif method == 1:
-            coordTransform = QgsCoordinateTransform(layer.crs(), context.project().crs())
+            coordTransform = QgsCoordinateTransform(source.sourceCrs(), context.project().crs())
 
-        features = QgsProcessingUtils.getFeatures(layer, context)
-        total = 100.0 / layer.featureCount() if layer.featureCount() else 0
+        features = source.getFeatures()
+        total = 100.0 / source.featureCount() if source.featureCount() else 0
         for current, f in enumerate(features):
+            if feedback.isCanceled():
+                break
+
             outFeat = f
             attrs = f.attributes()
             inGeom = f.geometry()
@@ -149,11 +149,11 @@ def processAlgorithm(self, parameters, context, feedback):
                     attrs.extend(self.line_attributes(inGeom))
 
             outFeat.setAttributes(attrs)
-            writer.addFeature(outFeat, QgsFeatureSink.FastInsert)
+            sink.addFeature(outFeat, QgsFeatureSink.FastInsert)
 
             feedback.setProgress(int(current * total))
 
-        del writer
+        return {self.OUTPUT: dest_id}
 
     def point_attributes(self, geometry):
         pt = None
diff --git a/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py b/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py
@@ -57,6 +57,7 @@
 from .Difference import Difference
 from .DropGeometry import DropGeometry
 from .DropMZValues import DropMZValues
+from .ExportGeometryInfo import ExportGeometryInfo
 from .ExtendLines import ExtendLines
 from .ExtentFromLayer import ExtentFromLayer
 from .ExtractNodes import ExtractNodes
@@ -117,7 +118,6 @@
 from .ZonalStatistics import ZonalStatistics
 
 # from .ExtractByLocation import ExtractByLocation
-# from .ExportGeometryInfo import ExportGeometryInfo
 # from .SinglePartsToMultiparts import SinglePartsToMultiparts
 # from .ConvexHull import ConvexHull
 # from .FixedDistanceBuffer import FixedDistanceBuffer
@@ -185,8 +185,6 @@ def __init__(self):
 
     def getAlgs(self):
         # algs = [
-        #
-        #         ExportGeometryInfo(),
         #         SinglePartsToMultiparts(),
         #         ConvexHull(), FixedDistanceBuffer(),
         #         VariableDistanceBuffer(),
@@ -244,6 +242,7 @@ def getAlgs(self):
                 Difference(),
                 DropGeometry(),
                 DropMZValues(),
+                ExportGeometryInfo(),
                 ExtendLines(),
                 ExtentFromLayer(),
                 ExtractNodes(),
diff --git a/python/plugins/processing/tests/testdata/qgis_algorithm_tests.yaml b/python/plugins/processing/tests/testdata/qgis_algorithm_tests.yaml
@@ -1095,17 +1095,17 @@ tests:
         name: expected/smoothed_lines_max_angle.gml
         type: vector
 
-#  - algorithm: qgis:exportaddgeometrycolumns
-#    name: Add Geometry PointZ
-#    params:
-#      CALC_METHOD: '0'
-#      INPUT:
-#        name: pointsz.gml
-#        type: vector
-#    results:
-#      OUTPUT:
-#        name: expected/add_geometry_pointz.gml
-#        type: vector
+  - algorithm: qgis:exportaddgeometrycolumns
+    name: Add Geometry PointZ
+    params:
+      CALC_METHOD: '0'
+      INPUT:
+        name: pointsz.gml
+        type: vector
+    results:
+      OUTPUT:
+        name: expected/add_geometry_pointz.gml
+        type: vector
 
   - algorithm: qgis:countpointsinpolygon
     name: Count points in polygon
