Add graphviz representation of the tree

stree/Splitter.py

@@ -145,6 +145,30 @@ class Snode:
             except IndexError:
                 self._class = None
 
+    def graph(self):
+        """
+        Return a string representing the node in graphviz format
+        """
+        output = ""
+        count_values = np.unique(self._y, return_counts=True)
+        if self.is_leaf():
+            output += (
+                f'N{id(self)} [shape=box style=filled label="'
+                f"class={self._class} belief={self._belief: .3f} "
+                f"impurity={self._impurity:.3f} "
+                f'classes/samples={count_values}"];\n'
+            )
+        else:
+            output += (
+                f'N{id(self)} [label="#features={len(self._features)} '
+                f'classes/samples={count_values}"];\n'
+            )
+            output += f'N{id(self)} -> N{id(self.get_up())} [label="Up"];\n'
+            output += (
+                f'N{id(self)} -> N{id(self.get_down())} [label="Down"];\n'
+            )
+        return output
+
     def __str__(self) -> str:
         count_values = np.unique(self._y, return_counts=True)
         if self.is_leaf():

stree/Strees.py

@@ -476,6 +476,20 @@ class Stree(BaseEstimator, ClassifierMixin):
             tree = None
         return Siterator(tree)
 
+    def graph(self) -> str:
+        """Graphviz code representing the tree
+
+        Returns
+        -------
+        str
+            graphviz code
+        """
+        output = "digraph STree {\n"
+        for node in self:
+            output += node.graph()
+        output += "}\n"
+        return output
+
     def __str__(self) -> str:
         """String representation of the tree
 

stree/tests/Stree_test.py

@@ -666,3 +666,20 @@ class Stree_test(unittest.TestCase):
     def test_version(self):
         clf = Stree()
         self.assertEqual(__version__, clf.version())
+
+    def test_graph(self):
+        X, y = load_wine(return_X_y=True)
+        clf = Stree(random_state=self._random_state)
+        clf.fit(X, y)
+        expected_head = "digraph STree {\n"
+        expected_tail = (
+            ' [shape=box style=filled label="class=1 belief= '
+            '1.000 impurity=0.000 classes/samples=(array([1]), array([1]))"]'
+            ";\n}\n"
+        )
+        computed = clf.graph()
+        computed_head = computed[: len(expected_head)]
+        num = -len(expected_tail)
+        computed_tail = computed[num:]
+        self.assertEqual(computed_head, expected_head)
+        self.assertEqual(computed_tail, expected_tail)