def transform()

in eland/ml/transformers/sklearn.py [0:0]

• 21 lines of code
• 7 McCabe index (conditional complexity)

    def transform(self) -> Tree:
        """
        Transform the provided model into an ES supported Tree object
        :return: Tree object for ES storage and use
        """
        target_type = (
            "regression"
            if isinstance(self._model, DecisionTreeRegressor)
            else "classification"
        )
        check_is_fitted(self._model, ["tree_"])
        tree_classes = None
        if self._classification_labels:
            tree_classes = self._classification_labels
        if isinstance(self._model, DecisionTreeClassifier):
            check_is_fitted(self._model, ["classes_"])
            if tree_classes is None:
                tree_classes = [str(c) for c in self._model.classes_]
        nodes = []
        tree_state = self._model.tree_.__getstate__()
        for i in range(len(tree_state["nodes"])):
            nodes.append(
                self.build_tree_node(i, tree_state["nodes"][i], tree_state["values"][i])
            )

        return Tree(self._feature_names, target_type, nodes, tree_classes)