Make memory layer indicator respond correctly to making a layer permanent

nyalldawson · nyalldawson · commit 01876d157189 · 2018-08-06T13:51:49.000+10:00
diff --git a/src/app/qgslayertreeviewmemoryindicator.cpp b/src/app/qgslayertreeviewmemoryindicator.cpp
@@ -17,6 +17,7 @@
 #include "qgslayertree.h"
 #include "qgslayertreemodel.h"
 #include "qgslayertreeview.h"
+#include "qgsvectorlayer.h"
 
 QgsLayerTreeViewMemoryIndicatorProvider::QgsLayerTreeViewMemoryIndicatorProvider( QgsLayerTreeView *view )
   : QObject( view )
@@ -28,6 +29,7 @@ QgsLayerTreeViewMemoryIndicatorProvider::QgsLayerTreeViewMemoryIndicatorProvider
   onAddedChildren( tree, 0, tree->children().count() - 1 );
 
   connect( tree, &QgsLayerTree::addedChildren, this, &QgsLayerTreeViewMemoryIndicatorProvider::onAddedChildren );
+  connect( tree, &QgsLayerTree::willRemoveChildren, this, &QgsLayerTreeViewMemoryIndicatorProvider::onWillRemoveChildren );
 }
 
 void QgsLayerTreeViewMemoryIndicatorProvider::onAddedChildren( QgsLayerTreeNode *node, int indexFrom, int indexTo )
@@ -46,8 +48,11 @@ void QgsLayerTreeViewMemoryIndicatorProvider::onAddedChildren( QgsLayerTreeNode
     {
       if ( QgsLayerTreeLayer *layerNode = dynamic_cast< QgsLayerTreeLayer * >( childNode ) )
       {
-        if ( layerNode->layer() && layerNode->layer()->dataProvider()->name() == QLatin1String( "memory" ) )
-          addIndicatorForMemoryLayer( childNode );
+        if ( QgsVectorLayer *vlayer = qobject_cast<QgsVectorLayer *>( layerNode->layer() ) )
+        {
+          connect( vlayer, &QgsVectorLayer::dataSourceChanged, this, &QgsLayerTreeViewMemoryIndicatorProvider::onDataSourceChanged );
+          addOrRemoveIndicator( childNode, vlayer );
+        }
         else if ( !layerNode->layer() )
         {
           // wait for layer to be loaded (e.g. when loading project, first the tree is loaded, afterwards the references to layers are resolved)
@@ -58,14 +63,63 @@ void QgsLayerTreeViewMemoryIndicatorProvider::onAddedChildren( QgsLayerTreeNode
   }
 }
 
+void QgsLayerTreeViewMemoryIndicatorProvider::onWillRemoveChildren( QgsLayerTreeNode *node, int indexFrom, int indexTo )
+{
+  // recursively disconnect from providers' dataChanged() signal
+
+  QList<QgsLayerTreeNode *> children = node->children();
+  for ( int i = indexFrom; i <= indexTo; ++i )
+  {
+    QgsLayerTreeNode *childNode = children[i];
+
+    if ( QgsLayerTree::isGroup( childNode ) )
+    {
+      onWillRemoveChildren( childNode, 0, childNode->children().count() - 1 );
+    }
+    else if ( QgsLayerTree::isLayer( childNode ) )
+    {
+      QgsLayerTreeLayer *childLayerNode = QgsLayerTree::toLayer( childNode );
+      if ( QgsVectorLayer *vlayer = qobject_cast<QgsVectorLayer *>( childLayerNode->layer() ) )
+      {
+        if ( vlayer )
+          disconnect( vlayer, &QgsVectorLayer::dataSourceChanged, this, &QgsLayerTreeViewMemoryIndicatorProvider::onDataSourceChanged );
+      }
+    }
+  }
+}
+
 void QgsLayerTreeViewMemoryIndicatorProvider::onLayerLoaded()
 {
   QgsLayerTreeLayer *layerNode = qobject_cast<QgsLayerTreeLayer *>( sender() );
   if ( !layerNode )
     return;
 
-  if ( layerNode->layer() && layerNode->layer()->dataProvider()->name() == QLatin1String( "memory" ) )
-    addIndicatorForMemoryLayer( layerNode );
+  if ( QgsVectorLayer *vlayer = qobject_cast<QgsVectorLayer *>( layerNode->layer() ) )
+  {
+    if ( vlayer )
+    {
+      connect( vlayer, &QgsVectorLayer::dataSourceChanged, this, &QgsLayerTreeViewMemoryIndicatorProvider::onDataSourceChanged );
+      addOrRemoveIndicator( layerNode, vlayer );
+    }
+  }
+}
+
+void QgsLayerTreeViewMemoryIndicatorProvider::onDataSourceChanged()
+{
+  QgsVectorLayer *vlayer = qobject_cast<QgsVectorLayer *>( sender() );
+  if ( !vlayer )
+    return;
+
+  // walk the tree and find layer node that needs to be updated
+  const QList<QgsLayerTreeLayer *> layerNodes = mLayerTreeView->layerTreeModel()->rootGroup()->findLayers();
+  for ( QgsLayerTreeLayer *node : layerNodes )
+  {
+    if ( node->layer() && node->layer() == vlayer )
+    {
+      addOrRemoveIndicator( node, vlayer );
+      break;
+    }
+  }
 }
 
 std::unique_ptr< QgsLayerTreeViewIndicator > QgsLayerTreeViewMemoryIndicatorProvider::newIndicator()
@@ -77,19 +131,41 @@ std::unique_ptr< QgsLayerTreeViewIndicator > QgsLayerTreeViewMemoryIndicatorProv
   return indicator;
 }
 
-void QgsLayerTreeViewMemoryIndicatorProvider::addIndicatorForMemoryLayer( QgsLayerTreeNode *node )
+void QgsLayerTreeViewMemoryIndicatorProvider::addOrRemoveIndicator( QgsLayerTreeNode *node, QgsVectorLayer *vlayer )
 {
-  const QList<QgsLayerTreeViewIndicator *> nodeIndicators = mLayerTreeView->indicators( node );
+  const bool isMemory = vlayer->dataProvider()->name() == QLatin1String( "memory" );
 
-  // maybe the indicator exists already
-  for ( QgsLayerTreeViewIndicator *indicator : nodeIndicators )
+  if ( isMemory )
   {
-    if ( mIndicators.contains( indicator ) )
+    const QList<QgsLayerTreeViewIndicator *> nodeIndicators = mLayerTreeView->indicators( node );
+
+    // maybe the indicator exists already
+    for ( QgsLayerTreeViewIndicator *indicator : nodeIndicators )
     {
-      return;
+      if ( mIndicators.contains( indicator ) )
+      {
+        return;
+      }
     }
+
+    // it does not exist: need to create a new one
+    mLayerTreeView->addIndicator( node, newIndicator().release() );
   }
+  else
+  {
+    const QList<QgsLayerTreeViewIndicator *> nodeIndicators = mLayerTreeView->indicators( node );
 
-  // it does not exist: need to create a new one
-  mLayerTreeView->addIndicator( node, newIndicator().release() );
+    // there may be existing indicator we need to get rid of
+    for ( QgsLayerTreeViewIndicator *indicator : nodeIndicators )
+    {
+      if ( mIndicators.contains( indicator ) )
+      {
+        mLayerTreeView->removeIndicator( node, indicator );
+        indicator->deleteLater();
+        return;
+      }
+    }
+
+    // no indicator was there before, nothing to do
+  }
 }
diff --git a/src/app/qgslayertreeviewmemoryindicator.h b/src/app/qgslayertreeviewmemoryindicator.h
@@ -23,6 +23,7 @@
 
 class QgsLayerTreeNode;
 class QgsLayerTreeView;
+class QgsVectorLayer;
 
 //! Adds indicators showing whether layers are memory layers.
 class QgsLayerTreeViewMemoryIndicatorProvider : public QObject
@@ -34,11 +35,15 @@ class QgsLayerTreeViewMemoryIndicatorProvider : public QObject
   private slots:
     //! Connects to signals of layers newly added to the tree
     void onAddedChildren( QgsLayerTreeNode *node, int indexFrom, int indexTo );
+    //! Disconnects from layers about to be removed from the tree
+    void onWillRemoveChildren( QgsLayerTreeNode *node, int indexFrom, int indexTo );
     void onLayerLoaded();
+    //! Adds/removes indicator of a layer
+    void onDataSourceChanged();
 
   private:
     std::unique_ptr< QgsLayerTreeViewIndicator > newIndicator();
-    void addIndicatorForMemoryLayer( QgsLayerTreeNode *node );
+    void addOrRemoveIndicator( QgsLayerTreeNode *node, QgsVectorLayer *vlayer );
 
   private:
     QgsLayerTreeView *mLayerTreeView = nullptr;
