Day 15 [WIP]

2 files changed, 14 insertions, 8 deletions

diff --git a/aoc/__init__.py b/aoc/__init__.py
index 8368f7b..7a089f6 100644
--- a/aoc/__init__.py
+++ b/aoc/__init__.py
@@ -1,11 +1,15 @@
 # -*- coding: utf-8 -*-
 import os
 from abc import ABC
-from typing import Generator, Any, Iterator
+from typing import Generator, Any, Iterator, Dict, TypeVar, Generic
 
+T = TypeVar("T")
+I = TypeVar("I")
 
-class BaseAssignment(ABC):
-    example_result = NotImplemented
+
+class BaseAssignment(Generic[T, I], ABC):
+    example_result: T = NotImplemented
+    example_kwargs: Dict = {}
 
     def __init__(self, path):
         self.path = path
@@ -13,14 +17,14 @@ class BaseAssignment(ABC):
     def __str__(self):
         return f"{self.__module__}.{self.__class__.__name__}"
 
-    def parse_item(self, item: str) -> Any:
+    def parse_item(self, item: str) -> I:
         return item
 
     @property
-    def part(self):
+    def part(self) -> int:
         return 1 if self.__class__.__name__.endswith("One") else 2
 
-    def read_input(self, example=False) -> Generator:
+    def read_input(self, example=False) -> Iterator[I]:
         file = f"{self.path}/input.txt"
 
         if example or not os.path.isfile(file):
@@ -35,5 +39,5 @@ class BaseAssignment(ABC):
             for line in input_file.readlines():
                 yield self.parse_item(line.strip("\n"))
 
-    def run(self, input: Iterator) -> Any:
+    def run(self, input: Iterator[I]) -> T:
         raise NotImplementedError("Please implement run")
diff --git a/aoc/test_init.py b/aoc/test_init.py
index f804c41..e52daf7 100644
--- a/aoc/test_init.py
+++ b/aoc/test_init.py
@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 def test_assignment_examples(assignment):
     assert (
-        assignment.run(input=assignment.read_input(example=True))
+        assignment.run(
+            input=assignment.read_input(example=True), **assignment.example_kwargs
+        )
         == assignment.example_result
     )