2 files changed, 37 insertions, 12 deletions

diff --git a/day9/__init__.py b/day9/__init__.py
index afa44aa..f198deb 100644
--- a/day9/__init__.py
+++ b/day9/__init__.py
@@ -2,7 +2,7 @@
 from abc import ABC
 from copy import copy
 from dataclasses import dataclass
-from typing import Tuple, Iterator, List
+from typing import Tuple, Iterator, List, Set
 
 from aoc import BaseAssignment
 
@@ -80,28 +80,39 @@ class Assignment(BaseAssignment, ABC):
                 self.next_tail(head, tail)
                 yield tail
 
-
-class AssignmentOne(Assignment):
-    example_result = 13
-
-    def run(self, input: Iterator) -> int:
+    def unique_tail_positions(self, input: Iterator[str]) -> Set[Coordinate]:
         unique_tail_positions = set()
 
         for position in self.tail_positions(input):
             unique_tail_positions.add(copy(position))
 
-        print(unique_tail_positions)
+        return unique_tail_positions
 
-        for y in range(5):
+    def visualize(self, width: int, height: int, positions: Set[Coordinate]):
+        rows = []
+        for y in range(height):
             items = []
-            for x in range(6):
-                if Coordinate(x, y) in unique_tail_positions:
+            for x in range(width):
+                if x == 0 and y == 0:
+                    items.append("s")
+                elif Coordinate(x, y) in positions:
                     items.append("#")
                 else:
                     items.append(".")
-            print("".join(items))
+            rows.append(items)
+
+        return "\n".join(["".join(row) for row in reversed(rows)])
+
+
+class AssignmentOne(Assignment):
+    example_result = 13
+
+    def run(self, input: Iterator) -> int:
+        unique_positions = self.unique_tail_positions(input)
+        print()
+        print(self.visualize(6, 5, unique_positions))
 
-        return len(unique_tail_positions)
+        return len(unique_positions)
 
 
 class AssignmentTwo(Assignment):
diff --git a/day9/test_init.py b/day9/test_init.py
new file mode 100644
index 0000000..b36eed6
--- /dev/null
+++ b/day9/test_init.py
@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+import os.path
+
+import day9
+
+
+def test_output_visualization():
+    expected = "..##..\n" "...##.\n" ".####.\n" "....#.\n" "s###.."
+
+    assignment = day9.AssignmentOne(path=os.path.dirname(day9.__file__))
+    unique_positions = assignment.unique_tail_positions(
+        input=assignment.read_input(True)
+    )
+    assert expected == assignment.visualize(6, 5, unique_positions)