Port Extract Nodes to new API

Improvements: - Retain Z/M values from input geometries
qgis · Jul 16, 2017 · 1b1dc7d · 1b1dc7d
1 parent bb344bc
commit 1b1dc7d
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 75 deletions.
diff --git a/python/plugins/processing/algs/qgis/ExtractNodes.py b/python/plugins/processing/algs/qgis/ExtractNodes.py
@@ -31,12 +31,16 @@
 from qgis.PyQt.QtGui import QIcon
 from qgis.PyQt.QtCore import QVariant
 
-from qgis.core import QgsFeature, QgsGeometry, QgsWkbTypes, QgsField, QgsFeatureSink, QgsProcessingUtils
+from qgis.core import (QgsFeature,
+                       QgsGeometry,
+                       QgsWkbTypes,
+                       QgsField,
+                       QgsFeatureSink,
+                       QgsProcessing,
+                       QgsProcessingParameterFeatureSource,
+                       QgsProcessingParameterFeatureSink)
 
 from processing.algs.qgis.QgisAlgorithm import QgisAlgorithm
-from processing.core.parameters import ParameterVector
-from processing.core.outputs import OutputVector
-from processing.tools import dataobjects, vector
 
 pluginPath = os.path.split(os.path.split(os.path.dirname(__file__))[0])[0]
 
@@ -56,12 +60,10 @@ def __init__(self):
         super().__init__()
 
     def initAlgorithm(self, config=None):
-        self.addParameter(ParameterVector(self.INPUT,
-                                          self.tr('Input layer'),
-                                          [dataobjects.TYPE_VECTOR_POLYGON,
-                                           dataobjects.TYPE_VECTOR_LINE]))
+        self.addParameter(QgsProcessingParameterFeatureSource(self.INPUT,
+                                                              self.tr('Input layer')))
 
-        self.addOutput(OutputVector(self.OUTPUT, self.tr('Nodes'), datatype=[dataobjects.TYPE_VECTOR_POINT]))
+        self.addParameter(QgsProcessingParameterFeatureSink(self.OUTPUT, self.tr('Nodes'), QgsProcessing.TypeVectorPoint))
 
     def name(self):
         return 'extractnodes'
@@ -70,36 +72,51 @@ def displayName(self):
         return self.tr('Extract nodes')
 
     def processAlgorithm(self, parameters, context, feedback):
-        layer = QgsProcessingUtils.mapLayerFromString(self.getParameterValue(self.INPUT), context)
+        source = self.parameterAsSource(parameters, self.INPUT, context)
 
-        fields = layer.fields()
+        fields = source.fields()
         fields.append(QgsField('node_index', QVariant.Int))
         fields.append(QgsField('distance', QVariant.Double))
         fields.append(QgsField('angle', QVariant.Double))
 
-        writer = self.getOutputFromName(self.OUTPUT).getVectorWriter(fields, QgsWkbTypes.Point, layer.crs(), context)
+        out_wkb = QgsWkbTypes.Point
+        if QgsWkbTypes.hasM(source.wkbType()):
+            out_wkb = QgsWkbTypes.addM(out_wkb)
+        if QgsWkbTypes.hasZ(source.wkbType()):
+            out_wkb = QgsWkbTypes.addZ(out_wkb)
 
-        features = QgsProcessingUtils.getFeatures(layer, context)
-        total = 100.0 / layer.featureCount() if layer.featureCount() else 0
+        (sink, dest_id) = self.parameterAsSink(parameters, self.OUTPUT, context,
+                                               fields, out_wkb, source.sourceCrs())
+
+        features = source.getFeatures()
+        total = 100.0 / source.featureCount() if source.featureCount() else 0
         for current, f in enumerate(features):
+            if feedback.isCanceled():
+                break
+
             input_geometry = f.geometry()
             if not input_geometry:
-                writer.addFeature(f, QgsFeatureSink.FastInsert)
+                sink.addFeature(f, QgsFeatureSink.FastInsert)
             else:
-                points = vector.extractPoints(input_geometry)
-
-                for i, point in enumerate(points):
-                    distance = input_geometry.distanceToVertex(i)
-                    angle = math.degrees(input_geometry.angleAtVertex(i))
-                    attrs = f.attributes()
-                    attrs.append(i)
-                    attrs.append(distance)
-                    attrs.append(angle)
-                    output_feature = QgsFeature()
-                    output_feature.setAttributes(attrs)
-                    output_feature.setGeometry(QgsGeometry.fromPoint(point))
-                    writer.addFeature(output_feature, QgsFeatureSink.FastInsert)
+                i = 0
+                for part in input_geometry.geometry().coordinateSequence():
+                    for ring in part:
+                        if feedback.isCanceled():
+                            break
+
+                        for point in ring:
+                            distance = input_geometry.distanceToVertex(i)
+                            angle = math.degrees(input_geometry.angleAtVertex(i))
+                            attrs = f.attributes()
+                            attrs.append(i)
+                            attrs.append(distance)
+                            attrs.append(angle)
+                            output_feature = QgsFeature()
+                            output_feature.setAttributes(attrs)
+                            output_feature.setGeometry(QgsGeometry(point.clone()))
+                            sink.addFeature(output_feature, QgsFeatureSink.FastInsert)
+                            i += 1
 
             feedback.setProgress(int(current * total))
 
-        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
@@ -57,6 +57,7 @@
 from .Difference import Difference
 from .DropGeometry import DropGeometry
 from .ExtentFromLayer import ExtentFromLayer
+from .ExtractNodes import ExtractNodes
 from .FixGeometry import FixGeometry
 from .GridPolygon import GridPolygon
 from .Heatmap import Heatmap
@@ -97,7 +98,6 @@
 # from .ExtractByLocation import ExtractByLocation
 # from .ExportGeometryInfo import ExportGeometryInfo
 # from .SinglePartsToMultiparts import SinglePartsToMultiparts
-# from .ExtractNodes import ExtractNodes
 # from .ConvexHull import ConvexHull
 # from .FixedDistanceBuffer import FixedDistanceBuffer
 # from .VariableDistanceBuffer import VariableDistanceBuffer
@@ -187,7 +187,6 @@ def getAlgs(self):
         #
         #         ExportGeometryInfo(),
         #         SinglePartsToMultiparts(),
-        #         ExtractNodes(),
         #         ConvexHull(), FixedDistanceBuffer(),
         #         VariableDistanceBuffer(),
         #         RandomSelection(), RandomSelectionWithinSubsets(),
@@ -251,6 +250,7 @@ def getAlgs(self):
                 Difference(),
                 DropGeometry(),
                 ExtentFromLayer(),
+                ExtractNodes(),
                 FixGeometry(),
                 GridPolygon(),
                 Heatmap(),

diff --git a/python/plugins/processing/tests/testdata/qgis_algorithm_tests.yaml b/python/plugins/processing/tests/testdata/qgis_algorithm_tests.yaml
@@ -889,50 +889,50 @@ tests:
 #      OUTPUT_LAYER:
 #        name: expected/single_sided_buffer_multiline_bevel.gml
 #        type: vector
-#
-#  - algorithm: qgis:extractnodes
-#    name: Test (qgis:extractnodes)
-#    params:
-#      INPUT:
-#        name: multipolys.gml
-#        type: vector
-#    results:
-#      OUTPUT:
-#        name: expected/extract_nodes_multipolys.gml
-#        type: vector
-#
-#  - algorithm: qgis:extractnodes
-#    name: Extract nodes from polygons
-#    params:
-#      INPUT:
-#        name: polys.gml
-#        type: vector
-#    results:
-#      OUTPUT:
-#        name: expected/extract_nodes_polys.gml
-#        type: vector
-#
-#  - algorithm: qgis:extractnodes
-#    name: Extract nodes from multilines
-#    params:
-#      INPUT:
-#        name: multilines.gml
-#        type: vector
-#    results:
-#      OUTPUT:
-#        name: expected/extract_nodes_multilines.gml
-#        type: vector
-#
-#  - algorithm: qgis:extractnodes
-#    name: Extract nodes from lines
-#    params:
-#      INPUT:
-#        name: lines.gml
-#        type: vector
-#    results:
-#      OUTPUT:
-#        name: expected/extract_nodes_lines.gml
-#        type: vector
+
+  - algorithm: qgis:extractnodes
+    name: Test (qgis:extractnodes)
+    params:
+      INPUT:
+        name: multipolys.gml
+        type: vector
+    results:
+      OUTPUT:
+        name: expected/extract_nodes_multipolys.gml
+        type: vector
+
+  - algorithm: qgis:extractnodes
+    name: Extract nodes from polygons
+    params:
+      INPUT:
+        name: polys.gml
+        type: vector
+    results:
+      OUTPUT:
+        name: expected/extract_nodes_polys.gml
+        type: vector
+
+  - algorithm: qgis:extractnodes
+    name: Extract nodes from multilines
+    params:
+      INPUT:
+        name: multilines.gml
+        type: vector
+    results:
+      OUTPUT:
+        name: expected/extract_nodes_multilines.gml
+        type: vector
+
+  - algorithm: qgis:extractnodes
+    name: Extract nodes from lines
+    params:
+      INPUT:
+        name: lines.gml
+        type: vector
+    results:
+      OUTPUT:
+        name: expected/extract_nodes_lines.gml
+        type: vector
 
   - algorithm: qgis:simplifygeometries
     name: Simplify (lines)