Add data roles for other columns

1 files changed, 20 insertions, 0 deletions

diff --git a/sigmodr/tree/TreeModel.cpp b/sigmodr/tree/TreeModel.cpp
index 5db31e06..0ec202fc 100644
--- a/sigmodr/tree/TreeModel.cpp
+++ b/sigmodr/tree/TreeModel.cpp
@@ -44,6 +44,26 @@ QVariant TreeModel::data(const QModelIndex& index, const int role) const
         if (!index.column())
             return item->text();
     }
+    else if (role == Qt::DecorationRole)
+    {
+        switch (index.column())
+        {
+            case 1:
+                if (item->canEdit())
+                    return KIcon("document-edit");
+                break;
+            case 2:
+                if (item->canRemove())
+                    return KIcon("list-remove");
+                break;
+            case 3:
+                if (item->canAddTo())
+                    return KIcon("list-add");
+                break;
+            default:
+                break;
+        }
+    }
     return QVariant();
 }