refactor Extract by attribute alg

qgis · Oct 3, 2014 · 5d114a8 · 5d114a8
1 parent 5306cd8
commit 5d114a8
Show file tree

Hide file tree

Showing 2 changed files with 126 additions and 1 deletion.
diff --git a/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py b/python/plugins/processing/algs/qgis/QGISAlgorithmProvider.py
@@ -76,6 +76,7 @@
 from mmqgisx.DeleteColumn import DeleteColumn
 from mmqgisx.DeleteDuplicateGeometries import DeleteDuplicateGeometries
 from mmqgisx.TextToFloat import TextToFloat
+from mmqgisx.ExtractByAttribute import ExtractByAttribute
 
 from ConcaveHull import ConcaveHull
 from Polygonize import Polygonize
@@ -141,7 +142,7 @@ def __init__(self):
                         SpatialJoin(),
                         # ------ mmqgisx ------
                         DeleteColumn(), DeleteDuplicateGeometries(),
-                        TextToFloat(),
+                        TextToFloat(), ExtractByAttribute(),
                         #mmqgisx_delete_duplicate_geometries_algorithm(),
                         #mmqgisx_geometry_convert_algorithm(),
                         #mmqgisx_grid_algorithm(),

diff --git a/python/plugins/processing/algs/qgis/mmqgisx/ExtractByAttribute.py b/python/plugins/processing/algs/qgis/mmqgisx/ExtractByAttribute.py
@@ -0,0 +1,124 @@
+# -*- coding: utf-8 -*-
+
+"""
+***************************************************************************
+    DeleteColumn.py
+    ---------------------
+    Date                 : May 2010
+    Copyright            : (C) 2010 by Michael Minn
+    Email                : pyqgis at michaelminn dot com
+***************************************************************************
+*                                                                         *
+*   This program is free software; you can redistribute it and/or modify  *
+*   it under the terms of the GNU General Public License as published by  *
+*   the Free Software Foundation; either version 2 of the License, or     *
+*   (at your option) any later version.                                   *
+*                                                                         *
+***************************************************************************
+"""
+
+__author__ = 'Michael Minn'
+__date__ = 'May 2010'
+__copyright__ = '(C) 2010, Michael Minn'
+
+# This will get replaced with a git SHA1 when you do a git archive
+
+__revision__ = '$Format:%H$'
+
+from PyQt4.QtCore import *
+from qgis.core import *
+from processing.core.GeoAlgorithm import GeoAlgorithm
+from processing.core.GeoAlgorithmExecutionException import \
+        GeoAlgorithmExecutionException
+from processing.core.parameters import ParameterVector
+from processing.core.parameters import ParameterTableField
+from processing.core.parameters import ParameterSelection
+from processing.core.parameters import ParameterString
+from processing.core.outputs import OutputVector
+from processing.tools import dataobjects, vector
+
+
+class ExtractByAttribute(GeoAlgorithm):
+    INPUT = 'INPUT'
+    FIELD = 'FIELD'
+    OPERATOR = 'OPERATOR'
+    VALUE = 'VALUE'
+    OUTPUT = 'OUTPUT'
+
+    OPERATORS = ['=',
+                 '!=',
+                 '>',
+                 '>=',
+                 '<',
+                 '<=',
+                 'begins with',
+                 'contains'
+                ]
+
+    def defineCharacteristics(self):
+        self.name = 'Extract by attribute'
+        self.group = 'Vector selection tools'
+
+        self.addParameter(ParameterVector(
+            self.INPUT, 'Input Layer', [ParameterVector.VECTOR_TYPE_ANY]))
+        self.addParameter(ParameterTableField(
+            self.FIELD, 'Selection attribute', self.INPUT))
+        self.addParameter(ParameterSelection(
+            self.OPERATOR, 'Operator', self.OPERATORS))
+        self.addParameter(ParameterString(self.VALUE, 'Value'))
+
+        self.addOutput(OutputVector(self.OUTPUT, 'Output'))
+
+    def processAlgorithm(self, progress):
+        layer = dataobjects.getObjectFromUri(
+                self.getParameterValue(self.INPUT))
+        fieldName = self.getParameterValue(self.FIELD)
+        operator = self.OPERATORS[self.getParameterValue(self.OPERATOR)]
+        value = self.getParameterValue(self.VALUE)
+
+        fields = layer.pendingFields()
+        writer = self.getOutputFromName(self.OUTPUT).getVectorWriter(fields,
+            layer.wkbType(), layer.crs())
+
+        idx = layer.fieldNameIndex(fieldName)
+        fieldType = fields[idx].type()
+
+        if fieldType != QVariant.String and operator in self.OPERATORS[-2:]:
+            op = ''.join(['"%s", ' % o for o in self.OPERATORS[-2:]])
+            raise GeoAlgorithmExecutionException(
+                'Operators %s can be used only with string fields.' % op)
+
+        if fieldType in [QVariant.Int, QVariant.Double]:
+            progress.setInfo('Numeric field')
+            expr = '"%s" %s %s' % (fieldName, operator, value)
+            progress.setInfo(expr)
+        elif fieldType == QVariant.String:
+            progress.setInfo('String field')
+            if operator not in self.OPERATORS[-2:]:
+                expr = """"%s" %s '%s'""" % (fieldName, operator, value)
+            elif operator == 'begins with':
+                expr = """"%s" LIKE '%s%%'""" % (fieldName, value)
+            elif operator == 'contains':
+                expr = """"%s" LIKE '%%%s%%'""" % (fieldName, value)
+            progress.setInfo(expr)
+        elif fieldType in [QVariant.Date, QVariant.DateTime]:
+            progress.setInfo('Date field')
+            expr = """"%s" %s '%s'""" % (fieldX, operator, value)
+            progress.setInfo(expr)
+        else:
+            raise GeoAlgorithmExecutionException(
+                'Unsupported field type "%s"' % fields[idx].typeName())
+
+        expression = QgsExpression(expr)
+        expression.prepare(fields)
+
+        features = vector.features(layer)
+
+        count = len(features)
+        total = 100.0 / float(count)
+        for count, f in enumerate(features):
+            if expression.evaluate(f, fields):
+                writer.addFeature(f)
+            progress.setPercentage(int(count * total))
+
+        del writer