Clean up

cov/__main__.py

@@ -1,13 +1,13 @@
 import sys
 import pytest
-from cov.utils import results
+import cov.utils
 
 
 def main() -> None:
     arguments = sys.argv
 
     pytest.main(arguments[1:])
-    print(results)
+    cov.utils.print_results()
 
 
 if __name__ == "__main__":

cov/benchmark.py

@@ -0,0 +1,13 @@
+from typing import List
+
+
+class Benchmark:
+    def __init__(self, amount_of_branches: int) -> None:
+        self.amount_of_branches: int = amount_of_branches
+        self.reached_branches: List[bool] = amount_of_branches * [False]
+
+    def mark(self, branch_id: int) -> None:
+        self.reached_branches[branch_id] = True
+
+    def is_valid(self, branch_id: int) -> bool:
+        return 0 <= branch_id < self.amount_of_branches

cov/utils.py

@@ -1,15 +1,9 @@
-from typing import Set, Dict, Callable
-from inspect import stack
+from typing import Dict, Callable
+from cov.benchmark import Benchmark
+import inspect
 
 __all__ = ["test", "mark"]
 
-
-class Benchmark:
-    def __init__(self, amount_of_branches: int) -> None:
-        self.amount_of_branches: int = amount_of_branches
-        self.reached: Set[int] = set()
-
-
 results: Dict[str, Benchmark] = {}
 
 
@@ -29,9 +23,14 @@ def wrapper(function: Callable) -> Callable:
 
 
 def mark(branch_id: int) -> None:
-    function_name = stack()[1].function
+    function_name = inspect.stack()[1].function
+    benchmark = results.get(function_name)
 
-    if function_name not in results:
+    if benchmark is None or benchmark.is_valid(branch_id) is False:
         raise KeyError()
 
-    results[function_name].reached.add(branch_id)
+    benchmark.mark(branch_id)
+
+
+def print_results() -> None:
+    print(results)