misc: minor refactor

.gitignore

@@ -1,4 +1,6 @@
 generated_models/
+test_sets/
+benchmarks/
 
 # Created by https://www.toptal.com/developers/gitignore/api/pycharm,python,tex
 # Edit at https://www.toptal.com/developers/gitignore?templates=pycharm,python,tex

src/generator/generator.py

@@ -11,16 +11,16 @@ def generate_model(self, level: Level) -> str:
         pass
 
 
-def _multi_range(start: int, end: int, offset: list[int], valid: set[int]):
+def _multi_range(start: int, end: int, offset: list[int], valid: set[int]) -> list:
     return [x for x in zip(range(start, end), *[range(start + o, end + o) for o in offset]) if set(x).issubset(valid)]
 
 
-def _move_bounds(level: Level, offset: int):
+def _move_bounds(level: Level, offset: int) -> (int, int):
     if offset < 0:
         return max(-offset, level.first_pos), level.last_pos + 1
     else:
         return level.first_pos, min(level.size - offset, level.last_pos + 1)
 
 
-def _flatten(list_2d):
+def _flatten(list_2d: list[list]) -> list:
     return list(itertools.chain.from_iterable(list_2d))

src/generator/jani_generators.py

@@ -8,7 +8,7 @@
 Expr = Identifier | dict | Number | bool
 
 
-def _binary_op(op: str, left: Expr, right: Expr):
+def _binary_op(op: str, left: Expr, right: Expr) -> Expr:
     return {
         "op": op,
         "left": left,

src/generator/prism_generators.py

@@ -102,7 +102,7 @@ def generate_model(self, level: Level) -> str:
 
     @staticmethod
     def _generate_board(level: Level) -> str:
-        def to_variable(name: str, value: bool):
+        def to_variable(name: str, value: bool) -> str:
             return f"{name}: bool init {str(value).lower()};"
 
         return '\n'.join(to_variable(f"box_{i}", level.board[i] == TileType.BOX) for i in
@@ -190,7 +190,7 @@ def generate_model(self, level: Level) -> str:
 
     @staticmethod
     def _generate_board(level: Level) -> str:
-        def to_variable(name: str, value: bool):
+        def to_variable(name: str, value: bool) -> str:
             return f"{name}: bool init {str(value).lower()};"
 
         return '\n'.join(to_variable(f"box_{i}", level.board[i] == TileType.BOX) for i in

src/parser/level.py

@@ -18,15 +18,15 @@ class Level(object):
     columns: int
 
     @property
-    def size(self):
+    def size(self) -> int:
         return self.rows * self.columns
 
     @cached_property
-    def first_pos(self):
+    def first_pos(self) -> int:
         return min(self.reachable_tiles)
 
     @cached_property
-    def last_pos(self):
+    def last_pos(self) -> int:
         return max(self.reachable_tiles)
 
     @cached_property
@@ -65,5 +65,5 @@ def _neighbors(pos: int) -> list[int]:
     def boxes(self) -> set[int]:
         return {i for i, t in enumerate(self.board) if t == TileType.BOX}
 
-    def is_wall(self, i: int):
+    def is_wall(self, i: int) -> bool:
         return self.board[i] == TileType.WALL

src/util/util.py

@@ -1,17 +1,17 @@
 import logging
 
+BYTE_UNITS = {
+    "B": 1024 ** 0,
+    "KB": 1024 ** 1,
+    "MB": 1024 ** 2,
+    "GB": 1024 ** 3,
+}
+
 
 def exit_with_error(message: str):
     logging.error(message)
     exit(-1)
 
 
-def convert_size(size, from_unit, to_unit):
-    units = {
-        "B": 1024 ** 0,
-        "KB": 1024 ** 1,
-        "MB": 1024 ** 2,
-        "GB": 1024 ** 3,
-    }
-
-    return size * units[from_unit] // units[to_unit]
+def convert_size(size: int, from_unit: str, to_unit: str) -> int:
+    return size * BYTE_UNITS[from_unit] // BYTE_UNITS[to_unit]