improve tree plotting

Author: xadupre

_unittests/ut_plotting/test_text_plot.py

@@ -65,6 +65,18 @@ def test_onnx_text_plot_tree_cls(self):
         self.assertIn("         T y=", res)
         self.assertIn("n_classes=3", res)
 
+    def test_onnx_text_plot_tree_cls_2(self):
+        iris = load_iris()
+        X_train, y_train = iris.data.astype(numpy.float32), iris.target
+        clr = DecisionTreeClassifier()
+        clr.fit(X_train, y_train)
+        model_def = to_onnx(
+            clr, X_train.astype(numpy.float32), options={"zipmap": False}
+        )
+        res = onnx_text_plot_tree(model_def.graph.node[0])
+        self.assertIn("n_classes=3", res)
+        print(res)
+
     @ignore_warnings((UserWarning, FutureWarning))
     def test_onnx_simple_text_plot_kmeans(self):
         x = numpy.random.randn(10, 3)

onnx_array_api/plotting/text_plot.py

@@ -114,9 +114,9 @@ def process_tree(atts, treeid):
         for i in range(len(short[f"{prefix}_treeids"])):
             idn = short[f"{prefix}_nodeids"][i]
             node = nodes[idn]
-            node.target_nodeids = idn
-            node.target_ids = short[f"{prefix}_ids"][i]
-            node.target_weights = short[f"{prefix}_weights"][i]
+            node.append_target(
+                id=short[f"{prefix}_ids"][i], weight=short[f"{prefix}_weights"][i]
+            )
 
         def iterate(nodes, node, depth=0, true_false=""):
             node.depth = depth