from sklearn.datasets import load_iris
from sklearn import tree

# Load iris dataset
iris = load_iris()

# Load decision tree classifiers: DecisionTree
clf_DT = tree.DecisionTreeClassifier(random_state=0)

# Build decision tree classifier from iris dataset
clf_DT = clf_DT.fit(iris.data, iris.target)

def export_text(
    decision_tree,
    *,
    feature_names=None,
    max_depth=10,
    spacing=3,
    decimals=2,
    show_weights=False,
)

feature_names = iris['feature_names']
r = tree.export_text(clf_DT, feature_names=feature_names, max_depth=2)
print(r)

class_names0 = None 
# Then: class_names = decision_tree.classes_

class_names1 = [0, "1"] 
# Then: class_names = decision_tree.classes_
# len(class_names1) != decision_tree.n_classes_

class_names2 = ["class0", "class1", "class2"] 
# Then: class_names = class_names2
# Reason: len(class_names2) == decision_tree.n_classes_

c_names = [class_names0, class_names1, class_names2]

print("===> After change <===")
for i, class_name in enumerate(c_names):
    r = tree.export_text(clf_DT, class_names=class_name, feature_names=feature_names, max_depth=2)
    print(r)

def export_text(
    decision_tree,
    *,
    feature_names=None,
    max_depth=10,
    spacing=3,
    decimals=2,
    show_weights=False,
):
    check_is_fitted(decision_tree)
    tree_ = decision_tree.tree_
    if is_classifier(decision_tree):
        class_names = decision_tree.classes_

def export_text(
    decision_tree,
    *,
    feature_names=None,
    class_names=None,  # add class_names argument
    max_depth=10,
    spacing=3,
    decimals=2,
    show_weights=False,
):
    """
    ...
    class_names : list of arguments, default=None
        Names of each of the target classes in ascending numerical order.
        Only relevant for classification and not supported for multi-output.
    ...
    """
    check_is_fitted(decision_tree)
    tree_ = decision_tree.tree_
    if is_classifier(decision_tree):
        if (  # length of class_names must be equal to the number of classes given by the tree
            class_names is not None
            and len(class_names) == len(decision_tree.classes_)
        ):
            class_names = class_names
        else:  # by default, we leave decision_tree.classes_
            class_names = decision_tree.classes_