diff --git a/ranked_tree.py b/ranked_tree.py index 92ded74..261c9ed 100644 --- a/ranked_tree.py +++ b/ranked_tree.py @@ -1,6 +1,6 @@ -from abc import abstractmethod from node import Node, Comparable +from abc import abstractmethod from collections import deque import logging from typing import Deque, Optional, Tuple, TypeVar, Generic @@ -13,12 +13,10 @@ class RankedTree(Generic[T]): def __init__(self) -> None: self.root: Optional[Node[T]] = None - def __str__(self) -> str: - result = "digraph {\n" - - queue: Deque[Optional[Node[T]]] = deque() - queue.append(self.root) + def _get_unwrapped_graph(self) -> str: + result = "" + queue: Deque[Optional[Node[T]]] = deque([self.root]) edges = [] while queue: node = queue.popleft() @@ -37,7 +35,10 @@ class RankedTree(Generic[T]): label = f'[label="{Node.difference(to_node)}"]' result += f'"{str(from_node)}" -> "{str(to_node)}" {label}\n' - return result + "}\n" + return result + + def __str__(self) -> str: + return "digraph {\n" + self._get_unwrapped_graph() + "}\n" # region TreeSpecificMethods