[FEATURE] new method QgsVectorLayer::selectByExpression(...)

nyalldawson · nyalldawson · commit b951d5a54d52 · 2016-05-19T09:11:00.000+10:00
Makes it simple for scripts to select by expression. The method
also accepts a parameter which dictates whether matching features
are added to an existing selection, removed from the selection
or intersected with the current selection.

The existing code from the select by expression dialog has been
moved to QgsVectorLayer, and optimised for maximum possible speed.

Also added unit tests.
diff --git a/python/core/qgsvectorlayer.sip b/python/core/qgsvectorlayer.sip
@@ -147,6 +147,15 @@ class QgsVectorLayer : QgsMapLayer
       InvalidLayer, /**< Edit failed due to invalid layer */
     };
 
+    //! Selection behaviour
+    enum SelectBehaviour
+    {
+      SetSelection, /**< Set selection, removing any existing selection */
+      AddToSelection, /**< Add selection to current selection */
+      IntersectSelection, /**< Modify current selection to include only select features which match */
+      RemoveFromSelection, /**<  Remove from current selection */
+    };
+
     /** Constructor - creates a vector layer
      *
      * The QgsVectorLayer is constructed by instantiating a data provider.  The provider
@@ -294,9 +303,20 @@ class QgsVectorLayer : QgsMapLayer
      * @param addToSelection  If set to true will not clear before selecting
      *
      * @see   invertSelectionInRectangle(QgsRectangle & rect)
+     * @see selectByExpression()
      */
     void select( QgsRectangle & rect, bool addToSelection );
 
+    /** Select matching features using an expression.
+     * @param expression expression to evaluate to select features
+     * @param behaviour selection type, allows adding to current selection, removing
+     * from selection, etc.
+     * @note added in QGIS 2.16
+     * @see select()
+     * @see modifySelection()
+     */
+    void selectByExpression( const QString& expression, SelectBehaviour behaviour = SetSelection );
+
     /**
      * Modifies the current selection on this layer
      *
@@ -307,6 +327,7 @@ class QgsVectorLayer : QgsMapLayer
      * @see   select(QgsFeatureId)
      * @see   deselect(QgsFeatureIds)
      * @see   deselect(QgsFeatureId)
+     * @see selectByExpression()
      */
     void modifySelection( QgsFeatureIds selectIds, QgsFeatureIds deselectIds );
 
diff --git a/src/core/qgsvectorlayer.cpp b/src/core/qgsvectorlayer.cpp
@@ -465,6 +465,71 @@ void QgsVectorLayer::select( QgsRectangle & rect, bool addToSelection )
   }
 }
 
+void QgsVectorLayer::selectByExpression( const QString& expression, QgsVectorLayer::SelectBehaviour behaviour )
+{
+  QgsFeatureIds newSelection;
+
+  QgsExpressionContext context;
+  context << QgsExpressionContextUtils::globalScope()
+  << QgsExpressionContextUtils::projectScope()
+  << QgsExpressionContextUtils::layerScope( this );
+
+  if ( behaviour == SetSelection || behaviour == AddToSelection )
+  {
+    QgsFeatureRequest request = QgsFeatureRequest().setFilterExpression( expression )
+                                .setExpressionContext( context )
+                                .setFlags( QgsFeatureRequest::NoGeometry );
+    //TODO - investigate whether removing all attributes is possible
+    //for now, just set the first attribute
+    request.setSubsetOfAttributes( QgsAttributeList() << 0 );
+
+    QgsFeatureIterator features = getFeatures( request );
+
+    if ( behaviour == AddToSelection )
+    {
+      newSelection = selectedFeaturesIds();
+    }
+    QgsFeature feat;
+    while ( features.nextFeature( feat ) )
+    {
+      newSelection << feat.id();
+    }
+    features.close();
+  }
+  else if ( behaviour == IntersectSelection || behaviour == RemoveFromSelection )
+  {
+    QgsExpression exp( expression );
+    exp.prepare( &context );
+
+    QgsFeatureIds oldSelection = selectedFeaturesIds();
+    QgsFeatureRequest request = QgsFeatureRequest().setFilterFids( oldSelection );
+
+    //refine request
+    if ( !exp.needsGeometry() )
+      request.setFlags( QgsFeatureRequest::NoGeometry );
+    request.setSubsetOfAttributes( exp.referencedColumns(), fields() );
+
+    QgsFeatureIterator features = getFeatures( request );
+    QgsFeature feat;
+    while ( features.nextFeature( feat ) )
+    {
+      context.setFeature( feat );
+      bool matches = exp.evaluate( &context ).toBool();
+
+      if ( matches && behaviour == IntersectSelection )
+      {
+        newSelection << feat.id();
+      }
+      else if ( !matches && behaviour == RemoveFromSelection )
+      {
+        newSelection << feat.id();
+      }
+    }
+  }
+
+  setSelectedFeatures( newSelection );
+}
+
 void QgsVectorLayer::modifySelection( QgsFeatureIds selectIds, QgsFeatureIds deselectIds )
 {
   QgsFeatureIds intersectingIds = selectIds & deselectIds;
diff --git a/src/core/qgsvectorlayer.h b/src/core/qgsvectorlayer.h
@@ -504,6 +504,15 @@ class CORE_EXPORT QgsVectorLayer : public QgsMapLayer
       InvalidLayer = 4, /**< Edit failed due to invalid layer */
     };
 
+    //! Selection behaviour
+    enum SelectBehaviour
+    {
+      SetSelection, /**< Set selection, removing any existing selection */
+      AddToSelection, /**< Add selection to current selection */
+      IntersectSelection, /**< Modify current selection to include only select features which match */
+      RemoveFromSelection, /**<  Remove from current selection */
+    };
+
     /** Constructor - creates a vector layer
      *
      * The QgsVectorLayer is constructed by instantiating a data provider.  The provider
@@ -657,9 +666,20 @@ class CORE_EXPORT QgsVectorLayer : public QgsMapLayer
      * @param addToSelection  If set to true will not clear before selecting
      *
      * @see   invertSelectionInRectangle(QgsRectangle & rect)
+     * @see selectByExpression()
      */
     void select( QgsRectangle & rect, bool addToSelection );
 
+    /** Select matching features using an expression.
+     * @param expression expression to evaluate to select features
+     * @param behaviour selection type, allows adding to current selection, removing
+     * from selection, etc.
+     * @note added in QGIS 2.16
+     * @see select()
+     * @see modifySelection()
+     */
+    void selectByExpression( const QString& expression, SelectBehaviour behaviour = SetSelection );
+
     /**
      * Modifies the current selection on this layer
      *
@@ -670,6 +690,7 @@ class CORE_EXPORT QgsVectorLayer : public QgsMapLayer
      * @see   select(QgsFeatureId)
      * @see   deselect(QgsFeatureIds)
      * @see   deselect(QgsFeatureId)
+     * @see selectByExpression()
      */
     void modifySelection( QgsFeatureIds selectIds, QgsFeatureIds deselectIds );
 
diff --git a/src/gui/qgsexpressionselectiondialog.cpp b/src/gui/qgsexpressionselectiondialog.cpp
@@ -76,138 +76,29 @@ void QgsExpressionSelectionDialog::setGeomCalculator( const QgsDistanceArea & da
 
 void QgsExpressionSelectionDialog::on_mActionSelect_triggered()
 {
-  QgsFeatureIds newSelection;
-  QgsExpression* expression = new QgsExpression( mExpressionBuilder->expressionText() );
-
-  QgsExpressionContext context;
-  context << QgsExpressionContextUtils::globalScope()
-  << QgsExpressionContextUtils::projectScope()
-  << QgsExpressionContextUtils::layerScope( mLayer );
-
-  QgsFeatureRequest request = QgsFeatureRequest().setFilterExpression( mExpressionBuilder->expressionText() ).setExpressionContext( context );
-  QgsFeatureIterator features = mLayer->getFeatures( request );
-
-  QgsFeature feat;
-  while ( features.nextFeature( feat ) )
-  {
-    newSelection << feat.id();
-  }
-
-  features.close();
-
-  mLayer->setSelectedFeatures( newSelection );
-
-  delete expression;
+  mLayer->selectByExpression( mExpressionBuilder->expressionText(),
+                              QgsVectorLayer::SetSelection );
   saveRecent();
 }
 
 void QgsExpressionSelectionDialog::on_mActionAddToSelection_triggered()
 {
-  QgsFeatureIds newSelection = mLayer->selectedFeaturesIds();
-  QgsExpression* expression = new QgsExpression( mExpressionBuilder->expressionText() );
-
-  QgsExpressionContext context;
-  context << QgsExpressionContextUtils::globalScope()
-  << QgsExpressionContextUtils::projectScope()
-  << QgsExpressionContextUtils::layerScope( mLayer );
-
-  QgsFeatureRequest request = QgsFeatureRequest().setFilterExpression( mExpressionBuilder->expressionText() ).setExpressionContext( context );
-  QgsFeatureIterator features = mLayer->getFeatures( request );
-
-  QgsFeature feat;
-  while ( features.nextFeature( feat ) )
-  {
-    newSelection << feat.id();
-  }
-
-  features.close();
-
-  mLayer->setSelectedFeatures( newSelection );
-
-  delete expression;
+  mLayer->selectByExpression( mExpressionBuilder->expressionText(),
+                              QgsVectorLayer::AddToSelection );
   saveRecent();
 }
 
 void QgsExpressionSelectionDialog::on_mActionSelectIntersect_triggered()
 {
-  const QgsFeatureIds &oldSelection = mLayer->selectedFeaturesIds();
-  QgsFeatureIds newSelection;
-
-  QgsExpression* expression = new QgsExpression( mExpressionBuilder->expressionText() );
-
-  QgsExpressionContext context;
-  context << QgsExpressionContextUtils::globalScope()
-  << QgsExpressionContextUtils::projectScope()
-  << QgsExpressionContextUtils::layerScope( mLayer );
-
-  expression->prepare( &context );
-
-  QgsFeature feat;
-  Q_FOREACH ( const QgsFeatureId fid, oldSelection )
-  {
-    QgsFeatureIterator features = mLayer->getFeatures( QgsFeatureRequest().setFilterFid( fid ) );
-
-    if ( features.nextFeature( feat ) )
-    {
-      context.setFeature( feat );
-      if ( expression->evaluate( &context ).toBool() )
-      {
-        newSelection << feat.id();
-      }
-    }
-    else
-    {
-      Q_ASSERT( false );
-    }
-
-    features.close();
-  }
-
-  mLayer->setSelectedFeatures( newSelection );
-
-  delete expression;
+  mLayer->selectByExpression( mExpressionBuilder->expressionText(),
+                              QgsVectorLayer::IntersectSelection );
   saveRecent();
 }
 
 void QgsExpressionSelectionDialog::on_mActionRemoveFromSelection_triggered()
 {
-  const QgsFeatureIds &oldSelection = mLayer->selectedFeaturesIds();
-  QgsFeatureIds newSelection = mLayer->selectedFeaturesIds();
-
-  QgsExpression* expression = new QgsExpression( mExpressionBuilder->expressionText() );
-
-  QgsExpressionContext context;
-  context << QgsExpressionContextUtils::globalScope()
-  << QgsExpressionContextUtils::projectScope()
-  << QgsExpressionContextUtils::layerScope( mLayer );
-
-  expression->prepare( &context );
-
-  QgsFeature feat;
-  Q_FOREACH ( const QgsFeatureId fid, oldSelection )
-  {
-    QgsFeatureIterator features = mLayer->getFeatures( QgsFeatureRequest().setFilterFid( fid ) );
-
-    if ( features.nextFeature( feat ) )
-    {
-      context.setFeature( feat );
-      if ( expression->evaluate( &context ).toBool() )
-      {
-        newSelection.remove( feat.id() );
-      }
-    }
-    else
-    {
-      Q_ASSERT( false );
-    }
-
-    features.close();
-  }
-
-  mLayer->setSelectedFeatures( newSelection );
-
-  delete expression;
-
+  mLayer->selectByExpression( mExpressionBuilder->expressionText(),
+                              QgsVectorLayer::RemoveFromSelection );
   saveRecent();
 }
 
diff --git a/tests/src/python/test_qgsvectorlayer.py b/tests/src/python/test_qgsvectorlayer.py
@@ -1063,6 +1063,38 @@ def test_ExpressionFilter(self):
 
         assert(len(list(features)) == 1)
 
+    def testSelectByExpression(self):
+        """ Test selecting by expression """
+        layer = QgsVectorLayer(os.path.join(unitTestDataPath(), 'points.shp'), 'Points', 'ogr')
+
+        # SetSelection
+        layer.selectByExpression('"Class"=\'B52\' and "Heading" > 10 and "Heading" <70', QgsVectorLayer.SetSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([10, 11]))
+        # check that existing selection is cleared
+        layer.selectByExpression('"Class"=\'Biplane\'', QgsVectorLayer.SetSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([1, 5, 6, 7, 8]))
+        # SelSelection no matching
+        layer.selectByExpression('"Class"=\'A380\'', QgsVectorLayer.SetSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([]))
+
+        # AddToSelection
+        layer.selectByExpression('"Importance"=3', QgsVectorLayer.AddToSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([0, 2, 3, 4, 14]))
+        layer.selectByExpression('"Importance"=4', QgsVectorLayer.AddToSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([0, 2, 3, 4, 13, 14]))
+
+        # IntersectSelection
+        layer.selectByExpression('"Heading"<100', QgsVectorLayer.IntersectSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([0, 2, 3, 4]))
+        layer.selectByExpression('"Cabin Crew"=1', QgsVectorLayer.IntersectSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([2, 3]))
+
+        # RemoveFromSelection
+        layer.selectByExpression('"Heading"=85', QgsVectorLayer.RemoveFromSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([3]))
+        layer.selectByExpression('"Heading"=95', QgsVectorLayer.RemoveFromSelection)
+        self.assertEqual(set(layer.selectedFeaturesIds()), set([]))
+
     def testAggregate(self):
         """ Test aggregate calculation """
         layer = QgsVectorLayer("Point?field=fldint:integer", "layer", "memory")
