Added HasTree marker interface

pmml-sklearn/src/main/java/sklearn2pmml/tree/CHAIDClassifier.java

@@ -28,7 +28,7 @@
 import sklearn.HasApplyField;
 import treelib.Tree;
 
-public class CHAIDClassifier extends Classifier implements HasApplyField {
+public class CHAIDClassifier extends Classifier implements HasApplyField, HasTree {
 
 	public CHAIDClassifier(String module, String name){
 		super(module, name);
@@ -41,18 +41,17 @@ public String getApplyField(){
 
 	@Override
 	public TreeModel encodeModel(Schema schema){
-		Tree tree = getTree();
-
 		CategoricalLabel categoricalLabel = (CategoricalLabel)schema.getLabel();
 
-		TreeModel treeModel = CHAIDUtil.encodeModel(MiningFunction.CLASSIFICATION, tree, schema);
+		TreeModel treeModel = CHAIDUtil.encodeModel(this, MiningFunction.CLASSIFICATION, schema);
 
 		encodePredictProbaOutput(treeModel, DataType.DOUBLE, categoricalLabel);
 		encodeApplyOutput(treeModel, DataType.INTEGER);
 
 		return treeModel;
 	}
 
+	@Override
 	public Tree getTree(){
 		return get("treelib_tree_", Tree.class);
 	}

pmml-sklearn/src/main/java/sklearn2pmml/tree/CHAIDRegressor.java

@@ -27,7 +27,7 @@
 import sklearn.Regressor;
 import treelib.Tree;
 
-public class CHAIDRegressor extends Regressor implements HasApplyField {
+public class CHAIDRegressor extends Regressor implements HasApplyField, HasTree {
 
 	public CHAIDRegressor(String module, String name){
 		super(module, name);
@@ -40,15 +40,14 @@ public String getApplyField(){
 
 	@Override
 	public TreeModel encodeModel(Schema schema){
-		Tree tree = getTree();
-
-		TreeModel treeModel = CHAIDUtil.encodeModel(MiningFunction.REGRESSION, tree, schema);
+		TreeModel treeModel = CHAIDUtil.encodeModel(this, MiningFunction.REGRESSION, schema);
 
 		encodeApplyOutput(treeModel, DataType.INTEGER);
 
 		return treeModel;
 	}
 
+	@Override
 	public Tree getTree(){
 		return get("treelib_tree_", Tree.class);
 	}

pmml-sklearn/src/main/java/sklearn2pmml/tree/CHAIDUtil.java

@@ -52,6 +52,7 @@
 import org.jpmml.converter.PredicateManager;
 import org.jpmml.converter.Schema;
 import org.jpmml.python.ClassDictUtil;
+import sklearn.Estimator;
 import treelib.Node;
 import treelib.Tree;
 
@@ -61,7 +62,9 @@ private CHAIDUtil(){
 	}
 
 	static
-	public TreeModel encodeModel(MiningFunction miningFunction, Tree tree, Schema schema){
+	public <E extends Estimator & HasTree> TreeModel encodeModel(E estimator, MiningFunction miningFunction, Schema schema){
+		Tree tree = estimator.getTree();
+
 		org.dmg.pmml.tree.Node root = encodeNode(True.INSTANCE, tree.selectRoot(), tree, new PredicateManager(), schema);
 
 		return new TreeModel(miningFunction, ModelUtil.createMiningSchema(schema.getLabel()), root);

pmml-sklearn/src/main/java/sklearn2pmml/tree/HasTree.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) 2024 Villu Ruusmann
+ *
+ * This file is part of JPMML-SkLearn
+ *
+ * JPMML-SkLearn is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * JPMML-SkLearn is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with JPMML-SkLearn.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package sklearn2pmml.tree;
+
+import treelib.Tree;
+
+public interface HasTree {
+
+	Tree getTree();
+}