OWTreeGraph: Output Flagged Data

Orange/tree.py

@@ -224,9 +224,14 @@ def _count(node):
         return _count(self.root)
 
     def get_instances(self, nodes):
+        indices = self.get_indices(nodes)
+        if indices is not None:
+            return self.instances[indices]
+
+    def get_indices(self, nodes):
         subsets = [node.subset for node in nodes]
         if subsets:
-            return self.instances[np.unique(np.hstack(subsets))]
+            return np.unique(np.hstack(subsets))
 
     @staticmethod
     def climb(node):

Orange/widgets/visualize/owtreeviewer.py

@@ -13,8 +13,10 @@
 
 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
+from Orange.widgets.utils.annotated_data import (create_annotated_table,
+                                                 ANNOTATED_DATA_SIGNAL_NAME)
 
 
 class PieChart(QGraphicsRectItem):
@@ -151,7 +153,8 @@ class OWTreeGraph(OWTreeViewer2D):
     icon = "icons/TreeViewer.svg"
     priority = 35
     inputs = [("Tree", TreeModel, "ctree")]
-    outputs = [("Data", Table)]
+    outputs = [("Selected Data", Table, widget.Default),
+               (ANNOTATED_DATA_SIGNAL_NAME, Table)]
 
     settingsHandler = ClassValuesContextHandler()
     target_class_index = ContextSetting(0)
@@ -266,7 +269,9 @@ 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(ANNOTATED_DATA_SIGNAL_NAME,
+                  create_annotated_table(self.dataset, None))
 
     def walkcreate(self, node_inst, parent=None):
         """Create a structure of tree nodes from the given model"""
@@ -291,7 +296,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(ANNOTATED_DATA_SIGNAL_NAME,
+                  create_annotated_table(self.dataset,
+                                         self.model.get_indices(nodes)))
 
     def send_report(self):
         if not self.model:

Orange/widgets/visualize/tests/test_owtreegraph.py

@@ -0,0 +1,28 @@
+# 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)
+        self.selected_indices = self.model.get_indices([node.node_inst])