From d65b52db7752e44a16366fd337d97c3b78999ca7 Mon Sep 17 00:00:00 2001
From: Nyall Dawson <nyall.dawson@gmail.com>
Date: Fri, 24 Jul 2020 08:58:07 +1000
Subject: [PATCH] Add CompatibleGeometryTypesRole to style model

---
 .../symbology/qgsstylemodel.sip.in            |  1 +
 src/core/symbology/qgsstylemodel.cpp          | 45 +++++++++++++++++--
 src/core/symbology/qgsstylemodel.h            |  1 +
 tests/src/python/test_qgsstylemodel.py        | 10 +++++
 4 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/python/core/auto_generated/symbology/qgsstylemodel.sip.in b/python/core/auto_generated/symbology/qgsstylemodel.sip.in
index 7784c700d11..15d0cf988a6 100644
--- a/python/core/auto_generated/symbology/qgsstylemodel.sip.in
+++ b/python/core/auto_generated/symbology/qgsstylemodel.sip.in
@@ -43,6 +43,7 @@ instead.
       SymbolTypeRole,
       IsFavoriteRole,
       LayerTypeRole,
+      CompatibleGeometryTypesRole,
     };
 
     explicit QgsStyleModel( QgsStyle *style, QObject *parent /TransferThis/ = 0 );
diff --git a/src/core/symbology/qgsstylemodel.cpp b/src/core/symbology/qgsstylemodel.cpp
index 9b87205a858..b8471d6b2b1 100644
--- a/src/core/symbology/qgsstylemodel.cpp
+++ b/src/core/symbology/qgsstylemodel.cpp
@@ -364,10 +364,49 @@ QVariant QgsStyleModel::data( const QModelIndex &index, int role ) const
 
     case LayerTypeRole:
     {
-      if ( entityType != QgsStyle::LabelSettingsEntity )
-        return QVariant();
+      switch ( entityType )
+      {
+        case QgsStyle::LabelSettingsEntity:
+          return mStyle->labelSettingsLayerType( name );
 
-      return mStyle->labelSettingsLayerType( name );
+        case QgsStyle::Symbol3DEntity:
+        case QgsStyle::SymbolEntity:
+        case QgsStyle::LegendPatchShapeEntity:
+        case QgsStyle::TagEntity:
+        case QgsStyle::ColorrampEntity:
+        case QgsStyle::SmartgroupEntity:
+        case QgsStyle::TextFormatEntity:
+          return QVariant();
+      }
+      return QVariant();
+    }
+
+    case CompatibleGeometryTypesRole:
+    {
+      switch ( entityType )
+      {
+        case QgsStyle::Symbol3DEntity:
+        {
+          QVariantList res;
+          const QList< QgsWkbTypes::GeometryType > types = mStyle->symbol3DCompatibleGeometryTypes( name );
+          res.reserve( types.size() );
+          for ( QgsWkbTypes::GeometryType type : types )
+          {
+            res << static_cast< int >( type );
+          }
+          return res;
+        }
+
+        case QgsStyle::LabelSettingsEntity:
+        case QgsStyle::SymbolEntity:
+        case QgsStyle::LegendPatchShapeEntity:
+        case QgsStyle::TagEntity:
+        case QgsStyle::ColorrampEntity:
+        case QgsStyle::SmartgroupEntity:
+        case QgsStyle::TextFormatEntity:
+          return QVariant();
+      }
+      return QVariant();
     }
 
     default:
diff --git a/src/core/symbology/qgsstylemodel.h b/src/core/symbology/qgsstylemodel.h
index 6ed0f2f4346..31d22404528 100644
--- a/src/core/symbology/qgsstylemodel.h
+++ b/src/core/symbology/qgsstylemodel.h
@@ -63,6 +63,7 @@ class CORE_EXPORT QgsStyleModel: public QAbstractItemModel
       SymbolTypeRole, //!< Symbol type (for symbol or legend patch shape entities)
       IsFavoriteRole, //!< Whether entity is flagged as a favorite
       LayerTypeRole, //!< Layer type (for label settings entities)
+      CompatibleGeometryTypesRole, //!< Compatible layer geometry types (for 3D symbols)
     };
 
     /**
diff --git a/tests/src/python/test_qgsstylemodel.py b/tests/src/python/test_qgsstylemodel.py
index 15a812465cf..9cd444394a1 100644
--- a/tests/src/python/test_qgsstylemodel.py
+++ b/tests/src/python/test_qgsstylemodel.py
@@ -54,6 +54,9 @@ class Dummy3dSymbol(QgsAbstract3DSymbol):
     def writeXml(self, elem, context):
         pass
 
+    def compatibleGeometryTypes(self):
+        return [int(QgsWkbTypes.PointGeometry), int(QgsWkbTypes.LineGeometry)]
+
 
 def createMarkerSymbol():
     symbol = QgsMarkerSymbol.createSimple({
@@ -496,6 +499,13 @@ class TestQgsStyleModel(unittest.TestCase):
         self.assertEqual(model.data(model.index(3, 0), QgsStyleModel.IsFavoriteRole), False)
         self.assertEqual(model.data(model.index(4, 0), QgsStyleModel.IsFavoriteRole), False)
 
+        self.assertEqual(model.data(model.index(0, 0), QgsStyleModel.CompatibleGeometryTypesRole), [0, 1])
+        self.assertEqual(model.data(model.index(1, 0), QgsStyleModel.CompatibleGeometryTypesRole), [0, 1])
+        self.assertEqual(model.data(model.index(2, 0), QgsStyleModel.CompatibleGeometryTypesRole), [0, 1])
+        self.assertEqual(model.data(model.index(3, 0), QgsStyleModel.CompatibleGeometryTypesRole), [0, 1])
+        self.assertEqual(model.data(model.index(4, 0), QgsStyleModel.CompatibleGeometryTypesRole), [0, 1])
+        self.assertEqual(model.data(model.index(5, 0), QgsStyleModel.CompatibleGeometryTypesRole), None)
+
     def test_mixed_style(self):
         """
         Test style with both symbols and ramps