Plotting decision trees with Graphviz with disk operations

graphviz_disk_op.py

import os
from typing import List

import graphviz
import matplotlib.image as plt_img
import numpy as np
from sklearn.tree import DecisionTreeClassifier, export_graphviz


def plot_disk_operations(clf: DecisionTreeClassifier,
                         feature_names: List[str],
                         class_names: List[str]) -> np.ndarray:
    # 1st disk operation: write DOT
    export_graphviz(clf, out_file="decision_tree.dot",
                    feature_names=feature_names,
                    class_names=class_names,
                    label="all", filled=True, impurity=False,
                    proportion=True, rounded=True, precision=2)

    # 2nd disk operation: read DOT
    graph = graphviz.Source.from_file("decision_tree.dot")

    # 3rd disk operation: write image
    graph.render("decision_tree", format="png")

    # 4th disk operation: read image
    image = plt_img.imread("decision_tree.png")

    # 5th and 6th disk operations: delete files
    os.remove("decision_tree.dot")
    os.remove("decision_tree.png")

    return image