Generate suffix tree of a given string

suffix_tree.py

from graphviz import Digraph
import argparse


class TreeNode:
    def __init__(self):
        self.children = {}
        self.indexes = []


class SuffixTree:
    def __init__(self, string):
        self.root = TreeNode()
        self.string = string
        self._construct_tree()

    def _construct_tree(self):
        for i in range(len(self.string)):
            current = self.root
            for char in self.string[i:]:
                if char not in current.children:
                    current.children[char] = TreeNode()
                current = current.children[char]
                current.indexes.append(i)

    def _draw_helper(self, node, dot, prefix):
        for char, child in node.children.items():
            child_name = f"{prefix}{char}"
            node_label = f"{char}"
            dot.node(child_name, node_label)
            dot.edge(prefix, child_name)
            if char == "$":
                dot.node(f"final{child_name}", f"{child.indexes[0]}")
                dot.edge(child_name, f"final{child_name}")
            self._draw_helper(child, dot, child_name)

    def save(self, filename):
        dot = Digraph()
        dot.node("root", ".")
        self._draw_helper(self.root, dot, "root")
        dot.render(filename, format="png", cleanup=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate a suffix tree for a given string."
    )
    parser.add_argument(
        "string", type=str, help="The string for which to generate the suffix tree."
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        default="suffix_tree",
        help="The output file name (without extension).",
    )

    args = parser.parse_args()
    string: str = args.string + "$"
    filename: str = args.output

    tree = SuffixTree(string)
    tree.save(filename)
    print(f"Suffix tree image generated as '{filename}.png'")