Port equivalent num field algorithm to new API

qgis · Jul 28, 2017 · 5b6696f · 5b6696f
1 parent fbff452
commit 5b6696f
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 43 deletions.
diff --git a/python/plugins/processing/algs/qgis/EquivalentNumField.py b/python/plugins/processing/algs/qgis/EquivalentNumField.py
@@ -27,14 +27,11 @@
 
 from qgis.PyQt.QtCore import QVariant
 from qgis.core import (QgsField,
-                       QgsFeature,
                        QgsFeatureSink,
-                       QgsApplication,
-                       QgsProcessingUtils)
+                       QgsProcessingParameterFeatureSource,
+                       QgsProcessingParameterField,
+                       QgsProcessingParameterFeatureSink)
 from processing.algs.qgis.QgisAlgorithm import QgisAlgorithm
-from processing.core.parameters import ParameterVector
-from processing.core.parameters import ParameterTableField
-from processing.core.outputs import OutputVector
 
 
 class EquivalentNumField(QgisAlgorithm):
@@ -50,11 +47,13 @@ def __init__(self):
         super().__init__()
 
     def initAlgorithm(self, config=None):
-        self.addParameter(ParameterVector(self.INPUT,
-                                          self.tr('Input layer')))
-        self.addParameter(ParameterTableField(self.FIELD,
-                                              self.tr('Class field'), self.INPUT))
-        self.addOutput(OutputVector(self.OUTPUT, self.tr('Layer with index field')))
+        self.addParameter(QgsProcessingParameterFeatureSource(self.INPUT,
+                                                              self.tr('Input layer')))
+        self.addParameter(QgsProcessingParameterField(self.FIELD,
+                                                      self.tr('Class field'),
+                                                      None, self.INPUT, QgsProcessingParameterField.Any))
+
+        self.addParameter(QgsProcessingParameterFeatureSink(self.OUTPUT, self.tr('Layer with index field')))
 
     def name(self):
         return 'adduniquevalueindexfield'
@@ -63,30 +62,34 @@ def displayName(self):
         return self.tr('Add unique value index field')
 
     def processAlgorithm(self, parameters, context, feedback):
-        fieldname = self.getParameterValue(self.FIELD)
-        output = self.getOutputFromName(self.OUTPUT)
-        vlayer = QgsProcessingUtils.mapLayerFromString(self.getParameterValue(self.INPUT), context)
-        fieldindex = vlayer.fields().lookupField(fieldname)
-        fields = vlayer.fields()
+        source = self.parameterAsSource(parameters, self.INPUT, context)
+        fields = source.fields()
         fields.append(QgsField('NUM_FIELD', QVariant.Int))
-        writer = output.getVectorWriter(fields, vlayer.wkbType(), vlayer.crs(), context)
-        outFeat = QgsFeature()
+
+        (sink, dest_id) = self.parameterAsSink(parameters, self.OUTPUT, context,
+                                               fields, source.wkbType(), source.sourceCrs())
+
+        field_name = self.parameterAsString(parameters, self.FIELD, context)
+        field_index = source.fields().lookupField(field_name)
+
         classes = {}
 
-        features = QgsProcessingUtils.getFeatures(vlayer, context)
-        total = 100.0 / vlayer.featureCount() if vlayer.featureCount() else 0
+        features = source.getFeatures()
+        total = 100.0 / source.featureCount() if source.featureCount() else 0
         for current, feature in enumerate(features):
+            if feedback.isCanceled():
+                break
+
             feedback.setProgress(int(current * total))
-            inGeom = feature.geometry()
-            outFeat.setGeometry(inGeom)
-            atMap = feature.attributes()
-            clazz = atMap[fieldindex]
+
+            attributes = feature.attributes()
+            clazz = attributes[field_index]
 
             if clazz not in classes:
                 classes[clazz] = len(list(classes.keys()))
 
-            atMap.append(classes[clazz])
-            outFeat.setAttributes(atMap)
-            writer.addFeature(outFeat, QgsFeatureSink.FastInsert)
+            attributes.append(classes[clazz])
+            feature.setAttributes(attributes)
+            sink.addFeature(feature, QgsFeatureSink.FastInsert)
 
-        del writer
+        return {self.OUTPUT: dest_id}
diff --git a/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py b/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py
@@ -58,6 +58,7 @@
 from .Difference import Difference
 from .DropGeometry import DropGeometry
 from .DropMZValues import DropMZValues
+from .EquivalentNumField import EquivalentNumField
 from .Explode import Explode
 from .ExportGeometryInfo import ExportGeometryInfo
 from .ExtendLines import ExtendLines
@@ -143,7 +144,6 @@
 # from .GeometryConvert import GeometryConvert
 # from .RasterLayerStatistics import RasterLayerStatistics
 # from .StatisticsByCategories import StatisticsByCategories
-# from .EquivalentNumField import EquivalentNumField
 # from .FieldsCalculator import FieldsCalculator
 # from .FieldPyculator import FieldsPyculator
 # from .JoinAttributes import JoinAttributes
@@ -198,7 +198,6 @@ def getAlgs(self):
         #         GeometryConvert(), FieldsCalculator(),
         #          JoinAttributes(),
         #         FieldsPyculator(),
-        #         EquivalentNumField(),
         #         StatisticsByCategories(),
         #         RasterLayerStatistics(), PointsDisplacement(),
         #         PointsFromPolygons(),
@@ -237,6 +236,7 @@ def getAlgs(self):
                 Difference(),
                 DropGeometry(),
                 DropMZValues(),
+                EquivalentNumField(),
                 Explode(),
                 ExportGeometryInfo(),
                 ExtendLines(),

diff --git a/python/plugins/processing/tests/testdata/qgis_algorithm_tests.yaml b/python/plugins/processing/tests/testdata/qgis_algorithm_tests.yaml
@@ -2125,19 +2125,19 @@ tests:
 #        compare:
 #          geometry:
 #            precision: 5
-#
-#  - algorithm: qgis:adduniquevalueindexfield
-#    name: add unique field based on another field
-#    params:
-#      FIELD: id2
-#      INPUT:
-#        name: points.gml
-#        type: vector
-#    results:
-#      OUTPUT:
-#        name: expected/add_unique_field.gml
-#        type: vector
-#
+
+  - algorithm: qgis:adduniquevalueindexfield
+    name: add unique field based on another field
+    params:
+      FIELD: id2
+      INPUT:
+        name: points.gml
+        type: vector
+    results:
+      OUTPUT:
+        name: expected/add_unique_field.gml
+        type: vector
+
   - algorithm: qgis:linestopolygons
     name: convert lines to polygon
     params: