1 files changed, 6 insertions, 1 deletions
diff --git a/Swift/QtUI/Roster/QtTreeWidget.h b/Swift/QtUI/Roster/QtTreeWidget.h
index 13b6d8e..15eb22f 100644
--- a/Swift/QtUI/Roster/QtTreeWidget.h
+++ b/Swift/QtUI/Roster/QtTreeWidget.h
@@ -10,6 +10,7 @@
 #include "Swift/QtUI/Roster/QtTreeWidget.h"
 #include "Swift/QtUI/Roster/RosterModel.h"
 #include "Swift/QtUI/Roster/RosterDelegate.h"
+#include "Swift/QtUI/ContextMenus/QtContextMenu.h"
 
 namespace Swift {
 
@@ -20,18 +21,22 @@ class QtTreeWidget : public QTreeView, public TreeWidget {
 		~QtTreeWidget();
 		void show();
 		QtTreeWidgetItem* getRoot();
+		void setContextMenu(QtContextMenu* contextMenu);
 	private slots:
 		void handleItemActivated(const QModelIndex&);
 		void handleModelItemExpanded(const QModelIndex&, bool expanded);
 		void handleExpanded(const QModelIndex&);
 		void handleCollapsed(const QModelIndex&);
 		void handleDataChanged(const QModelIndex& topLeft, const QModelIndex& bottomRight);
+	protected:
+		void contextMenuEvent(QContextMenuEvent* event);
+
 	private:
 		void drawBranches(QPainter*, const QRect&, const QModelIndex&) const;
 		RosterModel* model_;
 		RosterDelegate* delegate_;
 		QtTreeWidgetItem* treeRoot_;
-		
+		QtContextMenu* contextMenu_;
 };
 
 }