From a4c896edbf4104af3adfea05a2b58201f4dfd3c3 Mon Sep 17 00:00:00 2001
From: Matej Focko <mfocko@redhat.com>
Date: Thu, 3 Feb 2022 11:08:19 +0100
Subject: [PATCH] refactor(tree): split dot generation

Signed-off-by: Matej Focko <mfocko@redhat.com>
---
 ranked_tree.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

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