From ceb2a5e203910a3cb3196e39b91a427fe5e86ad4 Mon Sep 17 00:00:00 2001
From: Vesna Tanko <tankovesna@hotmail.com>
Date: Fri, 14 Oct 2016 16:00:51 +0200
Subject: [PATCH] OWTreeGraph: Output Flagged Data

---
 Orange/tree.py                                |  5 ++++
 Orange/widgets/visualize/owtreeviewer.py      | 14 +++++++---
 .../visualize/tests/test_owtreegraph.py       | 27 +++++++++++++++++++
 3 files changed, 42 insertions(+), 4 deletions(-)
 create mode 100644 Orange/widgets/visualize/tests/test_owtreegraph.py

diff --git a/Orange/tree.py b/Orange/tree.py
index 7fe154b7a6c..c84d074fc0c 100644
--- a/Orange/tree.py
+++ b/Orange/tree.py
@@ -228,6 +228,11 @@ def get_instances(self, nodes):
         if subsets:
             return self.instances[np.unique(np.hstack(subsets))]
 
+    def get_indices(self, nodes):
+        subsets = [node.subset for node in nodes]
+        if subsets:
+            return np.unique(np.hstack(subsets))
+
     @staticmethod
     def climb(node):
         while node:
diff --git a/Orange/widgets/visualize/owtreeviewer.py b/Orange/widgets/visualize/owtreeviewer.py
index 2f229339ded..069e40b7c7d 100644
--- a/Orange/widgets/visualize/owtreeviewer.py
+++ b/Orange/widgets/visualize/owtreeviewer.py
@@ -5,6 +5,7 @@
 from PyQt4.QtGui import QColor, QBrush, QPen, QFontMetrics, QStyle, \
     QSizePolicy, QGraphicsRectItem, QGraphicsTextItem, QLabel, QComboBox
 
+from Orange.misc.flagged_data import create_flagged_table, FLAGGED_SIGNAL_NAME
 from Orange.tree import TreeModel
 from Orange.widgets.visualize.owtreeviewer2d import \
     GraphicsNode, GraphicsEdge, OWTreeViewer2D
@@ -13,7 +14,7 @@
 
 from Orange.widgets.settings import ContextSetting, ClassValuesContextHandler, \
     Setting
-from Orange.widgets import gui
+from Orange.widgets import gui, widget
 from Orange.widgets.utils.colorpalette import ContinuousPaletteGenerator
 
 
@@ -151,7 +152,8 @@ class OWTreeGraph(OWTreeViewer2D):
     icon = "icons/TreeViewer.svg"
     priority = 35
     inputs = [("Tree", TreeModel, "ctree")]
-    outputs = [("Data", Table)]
+    outputs = [("Selected Data", Table, widget.Default),
+               (FLAGGED_SIGNAL_NAME, Table)]
 
     settingsHandler = ClassValuesContextHandler()
     target_class_index = ContextSetting(0)
@@ -266,7 +268,8 @@ def ctree(self, model=None):
             self.info.setText('{} nodes, {} leaves'.
                               format(model.node_count(), model.leaf_count()))
         self.setup_scene()
-        self.send("Data", None)
+        self.send("Selected Data", None)
+        self.send(FLAGGED_SIGNAL_NAME, create_flagged_table(self.dataset, None))
 
     def walkcreate(self, node_inst, parent=None):
         """Create a structure of tree nodes from the given model"""
@@ -291,7 +294,10 @@ def update_selection(self):
         nodes = [item.node_inst for item in self.scene.selectedItems()
                  if isinstance(item, TreeNode)]
         data = self.model.get_instances(nodes)
-        self.send("Data", data)
+        self.send("Selected Data", data)
+        self.send(FLAGGED_SIGNAL_NAME,
+                  create_flagged_table(self.dataset,
+                                       self.model.get_indices(nodes)))
 
     def send_report(self):
         if not self.model:
diff --git a/Orange/widgets/visualize/tests/test_owtreegraph.py b/Orange/widgets/visualize/tests/test_owtreegraph.py
new file mode 100644
index 00000000000..5dbbbc3e238
--- /dev/null
+++ b/Orange/widgets/visualize/tests/test_owtreegraph.py
@@ -0,0 +1,27 @@
+# Test methods with long descriptive names can omit docstrings
+# pylint: disable=missing-docstring
+from Orange.classification import TreeLearner
+from Orange.widgets.tests.base import WidgetTest, WidgetOutputsTestMixin
+from Orange.widgets.visualize.owtreeviewer import \
+    OWTreeGraph
+
+
+class TestOWTreeGraph(WidgetTest, WidgetOutputsTestMixin):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        WidgetOutputsTestMixin.init(cls)
+
+        tree = TreeLearner()
+        cls.model = tree(cls.data)
+        cls.model.instances = cls.data
+
+        cls.signal_name = "Tree"
+        cls.signal_data = cls.model
+
+    def setUp(self):
+        self.widget = self.create_widget(OWTreeGraph)
+
+    def _select_data(self):
+        node = self.widget.scene.nodes()[0]
+        node.setSelected(True)