Restore vector split algorithm

qgis · Jun 22, 2017 · 45eac05 · 45eac05
1 parent d8c44cb
commit 45eac05
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 23 deletions.
diff --git a/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py b/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py
@@ -60,8 +60,8 @@
 from .PostGISExecuteSQL import PostGISExecuteSQL
 from .RegularPoints import RegularPoints
 from .SymmetricalDifference import SymmetricalDifference
+from .VectorSplit import VectorSplit
 
-# from .VectorSplit import VectorSplit
 # from .RandomExtract import RandomExtract
 # from .RandomExtractWithinSubsets import RandomExtractWithinSubsets
 # from .ExtractByLocation import ExtractByLocation
@@ -201,7 +201,6 @@ def getAlgs(self):
         #         SelectByLocation(), RandomExtract(),
         #         RandomExtractWithinSubsets(), ExtractByLocation(),
         #         SpatialJoin(),
-        #         VectorSplit(),
         #         DeleteDuplicateGeometries(), TextToFloat(),
         #         SelectByAttribute(),
         #         GridLine(), Gridify(), HubDistancePoints(),
@@ -263,7 +262,8 @@ def getAlgs(self):
                 Merge(),
                 PostGISExecuteSQL(),
                 RegularPoints(),
-                SymmetricalDifference()
+                SymmetricalDifference(),
+                VectorSplit()
                 ]
 
         if hasPlotly:

diff --git a/python/plugins/processing/algs/qgis/VectorSplit.py b/python/plugins/processing/algs/qgis/VectorSplit.py
@@ -29,7 +29,13 @@
 import os
 
 from qgis.PyQt.QtGui import QIcon
-from qgis.core import QgsProcessingUtils
+from qgis.core import (QgsProcessingUtils,
+                       QgsProcessingParameterFeatureSource,
+                       QgsProcessingParameterField,
+                       QgsProcessingParameterFolderOutput,
+                       QgsProcessingOutputFolder,
+                       QgsExpression,
+                       QgsFeatureRequest)
 
 from processing.algs.qgis.QgisAlgorithm import QgisAlgorithm
 from processing.core.parameters import ParameterVector
@@ -55,11 +61,17 @@ def group(self):
 
     def __init__(self):
         super().__init__()
-        self.addParameter(ParameterVector(self.INPUT,
-                                          self.tr('Input layer')))
-        self.addParameter(ParameterTableField(self.FIELD,
-                                              self.tr('Unique ID field'), self.INPUT))
-        self.addOutput(OutputDirectory(self.OUTPUT, self.tr('Output directory')))
+
+        self.addParameter(QgsProcessingParameterFeatureSource(self.INPUT,
+                                                              self.tr('Input layer')))
+
+        self.addParameter(QgsProcessingParameterField(self.FIELD,
+                                                      self.tr('Unique ID field'), None, self.INPUT))
+
+        self.addParameter(QgsProcessingParameterFolderOutput(self.OUTPUT,
+                                                             self.tr('Output directory')))
+
+        self.addOutput(QgsProcessingOutputFolder(self.OUTPUT, self.tr('Output directory')))
 
     def name(self):
         return 'splitvectorlayer'
@@ -68,29 +80,42 @@ def displayName(self):
         return self.tr('Split vector layer')
 
     def processAlgorithm(self, parameters, context, feedback):
-        layer = QgsProcessingUtils.mapLayerFromString(self.getParameterValue(self.INPUT), context)
-        fieldName = self.getParameterValue(self.FIELD)
-        directory = self.getOutputValue(self.OUTPUT)
+        source = self.parameterAsSource(parameters, self.INPUT, context)
+        fieldName = self.parameterAsString(parameters, self.FIELD, context)
+        directory = self.parameterAsString(parameters, self.OUTPUT, context)
 
         mkdir(directory)
 
-        fieldIndex = layer.fields().lookupField(fieldName)
-        uniqueValues = QgsProcessingUtils.uniqueValues(layer, fieldIndex, context)
-        baseName = os.path.join(directory, '{0}_{1}'.format(layer.name(), fieldName))
+        fieldIndex = source.fields().lookupField(fieldName)
+        uniqueValues = source.uniqueValues(fieldIndex)
+        baseName = os.path.join(directory, '{0}'.format(fieldName))
 
-        fields = layer.fields()
-        crs = layer.crs()
-        geomType = layer.wkbType()
+        fields = source.fields()
+        crs = source.sourceCrs()
+        geomType = source.wkbType()
 
         total = 100.0 / len(uniqueValues)
 
         for current, i in enumerate(uniqueValues):
+            if feedback.isCanceled():
+                break
             fName = u'{0}_{1}.shp'.format(baseName, str(i).strip())
+            feedback.pushInfo(self.tr('Creating layer: {}').format(fName))
 
-            writer, dest = QgsProcessingUtils.createFeatureSink(fName, context, fields, geomType, crs)
-            for f in QgsProcessingUtils.getFeatures(layer, context):
-                if f[fieldName] == i:
-                    writer.addFeature(f)
-            del writer
+            sink, dest = QgsProcessingUtils.createFeatureSink(fName, context, fields, geomType, crs)
+
+            filter = '{} = {}'.format(QgsExpression.quotedColumnRef(fieldName), QgsExpression.quotedValue(i))
+            req = QgsFeatureRequest().setFilterExpression(filter)
+
+            count = 0
+            for f in source.getFeatures(req):
+                if feedback.isCanceled():
+                    break
+                sink.addFeature(f)
+                count += 1
+            feedback.pushInfo(self.tr('Added {} features to layer').format(count))
+            del sink
 
             feedback.setProgress(int(current * total))
+
+        return {self.OUTPUT: directory}