[FEATURE] Add with_variable function

This function sets a variable for parts of an expression, therefore avoiding to repeatedly recalculate complex values.
qgis · Jul 7, 2017 · efff5f1 · efff5f1
1 parent 9562536
commit efff5f1
Show file tree

Hide file tree

Showing 4 changed files with 142 additions and 25 deletions.
diff --git a/python/core/expression/qgsexpressionfunction.sip b/python/core/expression/qgsexpressionfunction.sip
@@ -249,6 +249,11 @@ The help text for the function.
  :rtype: QVariant
 %End
 
+    virtual QVariant run( QgsExpressionNode::NodeList *args, const QgsExpressionContext *context, QgsExpression *parent );
+%Docstring
+ :rtype: QVariant
+%End
+
     bool operator==( const QgsExpressionFunction &other ) const;
 
     virtual bool handlesNull() const;

diff --git a/src/core/expression/qgsexpressionfunction.cpp b/src/core/expression/qgsexpressionfunction.cpp
@@ -47,6 +47,34 @@ const QString QgsExpressionFunction::helpText() const
   return mHelpText.isEmpty() ? QgsExpression::helpText( mName ) : mHelpText;
 }
 
+QVariant QgsExpressionFunction::run( QgsExpressionNode::NodeList *args, const QgsExpressionContext *context, QgsExpression *parent )
+{
+  // evaluate arguments
+  QVariantList argValues;
+  if ( args )
+  {
+    Q_FOREACH ( QgsExpressionNode *n, args->list() )
+    {
+      QVariant v;
+      if ( lazyEval() )
+      {
+        // Pass in the node for the function to eval as it needs.
+        v = QVariant::fromValue( n );
+      }
+      else
+      {
+        v = n->eval( parent, context );
+        ENSURE_NO_EVAL_ERROR;
+        if ( QgsExpressionUtils::isNull( v ) && !handlesNull() )
+          return QVariant(); // all "normal" functions return NULL, when any QgsExpressionFunction::Parameter is NULL (so coalesce is abnormal)
+      }
+      argValues.append( v );
+    }
+  }
+
+  return func( argValues, context, parent );
+}
+
 bool QgsExpressionFunction::usesGeometry( const QgsExpressionNodeFunction *node ) const
 {
   Q_UNUSED( node )
@@ -4202,6 +4230,7 @@ const QList<QgsExpressionFunction *> &QgsExpression::Functions()
 
     sFunctions
         << new QgsStaticExpressionFunction( QStringLiteral( "env" ), 1, fcnEnvVar, QStringLiteral( "General" ), QString() )
+        << new QgsSetVariableExpressionFunction()
         << new QgsStaticExpressionFunction( QStringLiteral( "attribute" ), 2, fcnAttribute, QStringLiteral( "Record" ), QString(), false, QSet<QString>() << QgsFeatureRequest::ALL_ATTRIBUTES )
 
         // functions for arrays
@@ -4246,3 +4275,92 @@ const QList<QgsExpressionFunction *> &QgsExpression::Functions()
   }
   return sFunctions;
 }
+
+QgsSetVariableExpressionFunction::QgsSetVariableExpressionFunction()
+  : QgsExpressionFunction( "set_variable", 3, QCoreApplication::tr( "General" ), "help text TODODOODO" )
+{
+
+}
+
+bool QgsSetVariableExpressionFunction::isStatic( const QgsExpressionNodeFunction *node, QgsExpression *parent, const QgsExpressionContext *context ) const
+{
+  bool isStatic = false;
+
+  QgsExpressionNode::NodeList *args = node->args();
+
+  if ( args->count() < 3 )
+    return false;
+
+  if ( args->at( 0 )->isStatic( parent, context ) && args->at( 1 )->isStatic( parent, context ) )
+  {
+    QVariant name = args->at( 0 )->eval( parent, context );
+    QVariant value = args->at( 1 )->eval( parent, context );
+
+    QgsExpressionContextScope *scope = new QgsExpressionContextScope();
+    scope->setVariable( name.toString(), value );
+
+    QgsExpressionContext *updatedContext = const_cast<QgsExpressionContext *>( context );
+    updatedContext->appendScope( scope );
+
+    if ( args->at( 3 )->isStatic( parent, updatedContext ) )
+      isStatic = true;
+    updatedContext->popScope();
+  }
+
+  return false;
+}
+
+QVariant QgsSetVariableExpressionFunction::run( QgsExpressionNode::NodeList *args, const QgsExpressionContext *context, QgsExpression *parent )
+{
+  QVariant result;
+
+  if ( args->count() < 3 )
+    // error
+    return result;
+
+  QVariant name = args->at( 0 )->eval( parent, context );
+  QVariant value = args->at( 1 )->eval( parent, context );
+
+  QgsExpressionContextScope *scope = new QgsExpressionContextScope();
+  scope->setVariable( name.toString(), value );
+
+  QgsExpressionContext *updatedContext = const_cast<QgsExpressionContext *>( context );
+  updatedContext->appendScope( scope );
+  result = args->at( 2 )->eval( parent, updatedContext );
+  delete updatedContext->popScope();
+
+  return result;
+}
+
+QVariant QgsSetVariableExpressionFunction::func( const QVariantList &values, const QgsExpressionContext *context, QgsExpression *parent )
+{
+  // This is a dummy function, all the real handling is in run
+  Q_UNUSED( values )
+  Q_UNUSED( context )
+  Q_UNUSED( parent )
+
+  Q_ASSERT( false );
+  return QVariant();
+}
+
+bool QgsSetVariableExpressionFunction::prepare( const QgsExpressionNodeFunction *node, QgsExpression *parent, const QgsExpressionContext *context ) const
+{
+  QgsExpressionNode::NodeList *args = node->args();
+
+  if ( args->count() < 3 )
+    // error
+    return false;
+
+  QVariant name = args->at( 0 )->prepare( parent, context );
+  QVariant value = args->at( 1 )->prepare( parent, context );
+
+  QgsExpressionContextScope *scope = new QgsExpressionContextScope();
+  scope->setVariable( name.toString(), value );
+
+  QgsExpressionContext *updatedContext = const_cast<QgsExpressionContext *>( context );
+  updatedContext->appendScope( scope );
+  args->at( 2 )->prepare( parent, updatedContext );
+  delete updatedContext->popScope();
+
+  return true;
+}
diff --git a/src/core/expression/qgsexpressionfunction.h b/src/core/expression/qgsexpressionfunction.h
@@ -24,6 +24,7 @@
 
 #include "qgis.h"
 #include "qgis_core.h"
+#include "qgsexpressionnode.h"
 
 class QgsExpressionNodeFunction;
 class QgsExpression;
@@ -273,6 +274,8 @@ class CORE_EXPORT QgsExpressionFunction
      */
     virtual QVariant func( const QVariantList &values, const QgsExpressionContext *context, QgsExpression *parent ) = 0;
 
+    virtual QVariant run( QgsExpressionNode::NodeList *args, const QgsExpressionContext *context, QgsExpression *parent );
+
     bool operator==( const QgsExpressionFunction &other ) const;
 
     virtual bool handlesNull() const;
@@ -451,6 +454,21 @@ class QgsStaticExpressionFunction : public QgsExpressionFunction
     QSet<QString> mReferencedColumns;
     bool mIsStatic = false;
 };
+
+class QgsSetVariableExpressionFunction : public QgsExpressionFunction
+{
+  public:
+    QgsSetVariableExpressionFunction();
+
+    bool isStatic( const QgsExpressionNodeFunction *node, QgsExpression *parent, const QgsExpressionContext *context ) const override;
+
+    QVariant run( QgsExpressionNode::NodeList *args, const QgsExpressionContext *context, QgsExpression *parent ) override;
+
+    QVariant func( const QVariantList &values, const QgsExpressionContext *context, QgsExpression *parent ) override;
+
+    bool prepare( const QgsExpressionNodeFunction *node, QgsExpression *parent, const QgsExpressionContext *context ) const override;
+};
+
 #endif
 
 #endif // QGSEXPRESSIONFUNCTION_H
diff --git a/src/core/expression/qgsexpressionnodeimpl.cpp b/src/core/expression/qgsexpressionnodeimpl.cpp
@@ -842,31 +842,7 @@ QVariant QgsExpressionNodeFunction::evalNode( QgsExpression *parent, const QgsEx
   QString name = QgsExpression::QgsExpression::Functions()[mFnIndex]->name();
   QgsExpressionFunction *fd = context && context->hasFunction( name ) ? context->function( name ) : QgsExpression::QgsExpression::Functions()[mFnIndex];
 
-  // evaluate arguments
-  QVariantList argValues;
-  if ( mArgs )
-  {
-    Q_FOREACH ( QgsExpressionNode *n, mArgs->list() )
-    {
-      QVariant v;
-      if ( fd->lazyEval() )
-      {
-        // Pass in the node for the function to eval as it needs.
-        v = QVariant::fromValue( n );
-      }
-      else
-      {
-        v = n->eval( parent, context );
-        ENSURE_NO_EVAL_ERROR;
-        if ( QgsExpressionUtils::isNull( v ) && !fd->handlesNull() )
-          return QVariant(); // all "normal" functions return NULL, when any QgsExpressionFunction::Parameter is NULL (so coalesce is abnormal)
-      }
-      argValues.append( v );
-    }
-  }
-
-  // run the function
-  QVariant res = fd->func( argValues, context, parent );
+  QVariant res = fd->run( mArgs, context, parent );
   ENSURE_NO_EVAL_ERROR;
 
   // everything went fine