Amazing (Python): Code Cleanup

* Make more use of a Maze class
* Create a bigger maze in the test
* Group row/col variables in Position class
* Group direction variables in an Enum

02_Amazing/python/amazing.py

@@ -1,22 +1,86 @@
 import random
-from typing import List, NamedTuple, Tuple
+import enum
+from typing import List, Tuple
+from dataclasses import dataclass
 
 # Python translation by Frank Palazzolo - 2/2021
 
 
-class Maze(NamedTuple):
+class Maze:
-    used: List[List[int]]
+    def __init__(
-    walls: List[List[int]]
+        self,
-    enter_col: int
+        width: int,
-    width: int
+        length: int,
-    length: int
+    ):
+        assert width >= 2 and length >= 2
+        used: List[List[int]] = []
+        walls: List[List[int]] = []
+        for _ in range(length):
+            used.append([0] * width)
+            walls.append([0] * width)
+
+        # Pick a random entrance, mark as used
+        enter_col = random.randint(0, width - 1)
+        used[0][enter_col] = 1
+
+        self.used = used
+        self.walls = walls
+        self.enter_col = enter_col
+        self.width = width
+        self.length = length
+
+    def add_exit(self) -> None:
+        """Modifies 'walls' to add an exit to the maze."""
+        col = random.randint(0, self.width - 1)
+        row = self.length - 1
+        self.walls[row][col] = self.walls[row][col] + 1
+
+    def display(self) -> None:
+        for col in range(self.width):
+            if col == self.enter_col:
+                print(".  ", end="")
+            else:
+                print(".--", end="")
+        print(".")
+        for row in range(self.length):
+            print("I", end="")
+            for col in range(self.width):
+                if self.walls[row][col] < 2:
+                    print("  I", end="")
+                else:
+                    print("   ", end="")
+            print()
+            for col in range(self.width):
+                if self.walls[row][col] == 0 or self.walls[row][col] == 2:
+                    print(":--", end="")
+                else:
+                    print(":  ", end="")
+            print(".")
+
+
+class Direction(enum.Enum):
+    LEFT = 0
+    UP = 1
+    RIGHT = 2
+    DOWN = 3
+
+
+@dataclass
+class Position:
+    row: int
+    col: int
+
+
+# Give Exit directions nice names
+EXIT_DOWN = 1
+EXIT_RIGHT = 2
 
 
 def main() -> None:
     welcome_header()
     width, length = get_maze_dimensions()
     maze = build_maze(width, length)
-    print_maze(maze)
+    maze.display()
 
 
 def welcome_header() -> None:
@@ -28,7 +92,7 @@ def welcome_header() -> None:
 
 
 def build_maze(width: int, length: int) -> Maze:
-    # Build two 2D arrays
+    """Build two 2D arrays."""
     #
     # used:
     #   Initially set to zero, unprocessed cells
@@ -42,107 +106,89 @@ def build_maze(width: int, length: int) -> Maze:
     #   Set to 3 if there are exits down and right
     assert width >= 2 and length >= 2
 
-    used = []
+    maze = Maze(width, length)
-    walls = []
+    position = Position(row=0, col=maze.enter_col)
-    for _ in range(length):
+    count = 2
-        used.append([0] * width)
-        walls.append([0] * width)
-
-    # Use direction variables with nice names
-    GO_LEFT, GO_UP, GO_RIGHT, GO_DOWN = [0, 1, 2, 3]
-    # Give Exit directions nice names
-    EXIT_DOWN = 1
-    EXIT_RIGHT = 2
-
-    # Pick a random entrance, mark as used
-    enter_col = random.randint(0, width - 1)
-    row, col = 0, enter_col
-    count = 1
-    used[row][col] = count
-    count = count + 1
 
     while count != width * length + 1:
-        # remove possible directions that are blocked or
+        possible_dirs = get_possible_directions(maze, position)
-        # hit cells that we have already processed
-        possible_dirs = [GO_LEFT, GO_UP, GO_RIGHT, GO_DOWN]
-        if col == 0 or used[row][col - 1] != 0:
-            possible_dirs.remove(GO_LEFT)
-        if row == 0 or used[row - 1][col] != 0:
-            possible_dirs.remove(GO_UP)
-        if col == width - 1 or used[row][col + 1] != 0:
-            possible_dirs.remove(GO_RIGHT)
-        if row == length - 1 or used[row + 1][col] != 0:
-            possible_dirs.remove(GO_DOWN)
 
         # If we can move in a direction, move and make opening
         if len(possible_dirs) != 0:
-            direction = random.choice(possible_dirs)
+            position, count = make_opening(maze, possible_dirs, position, count)
-            if direction == GO_LEFT:
-                col = col - 1
-                walls[row][col] = EXIT_RIGHT
-            elif direction == GO_UP:
-                row = row - 1
-                walls[row][col] = EXIT_DOWN
-            elif direction == GO_RIGHT:
-                walls[row][col] = walls[row][col] + EXIT_RIGHT
-                col = col + 1
-            elif direction == GO_DOWN:
-                walls[row][col] = walls[row][col] + EXIT_DOWN
-                row = row + 1
-            used[row][col] = count
-            count = count + 1
         # otherwise, move to the next used cell, and try again
         else:
             while True:
-                if col != width - 1:
+                if position.col != width - 1:
-                    col = col + 1
+                    position.col = position.col + 1
-                elif row != length - 1:
+                elif position.row != length - 1:
-                    row, col = row + 1, 0
+                    position.row, position.col = position.row + 1, 0
                 else:
-                    row, col = 0, 0
+                    position.row, position.col = 0, 0
-                if used[row][col] != 0:
+                if maze.used[position.row][position.col] != 0:
                     break
 
-    # Add a random exit
+    maze.add_exit()
-    col = random.randint(0, width - 1)
+    return maze
-    row = length - 1
+
-    walls[row][col] = walls[row][col] + 1
+
-    return Maze(used, walls, enter_col, width, length)
+def make_opening(
+    maze: Maze,
+    possible_dirs: List[Direction],
+    pos: Position,
+    count: int,
+) -> Tuple[Position, int]:
+    """
+    Attention! This modifies 'used' and 'walls'
+    """
+    direction = random.choice(possible_dirs)
+    if direction == Direction.LEFT:
+        pos.col = pos.col - 1
+        maze.walls[pos.row][pos.col] = EXIT_RIGHT
+    elif direction == Direction.UP:
+        pos.row = pos.row - 1
+        maze.walls[pos.row][pos.col] = EXIT_DOWN
+    elif direction == Direction.RIGHT:
+        maze.walls[pos.row][pos.col] = maze.walls[pos.row][pos.col] + EXIT_RIGHT
+        pos.col = pos.col + 1
+    elif direction == Direction.DOWN:
+        maze.walls[pos.row][pos.col] = maze.walls[pos.row][pos.col] + EXIT_DOWN
+        pos.row = pos.row + 1
+    maze.used[pos.row][pos.col] = count
+    count = count + 1
+    return pos, count
+
+
+def get_possible_directions(maze: Maze, pos: Position) -> List[Direction]:
+    """
+    Get a list of all directions that are not blocked.
+
+    Also ignore hit cells that we have already processed
+    """
+    possible_dirs = list(Direction)
+    if pos.col == 0 or maze.used[pos.row][pos.col - 1] != 0:
+        possible_dirs.remove(Direction.LEFT)
+    if pos.row == 0 or maze.used[pos.row - 1][pos.col] != 0:
+        possible_dirs.remove(Direction.UP)
+    if pos.col == maze.width - 1 or maze.used[pos.row][pos.col + 1] != 0:
+        possible_dirs.remove(Direction.RIGHT)
+    if pos.row == maze.length - 1 or maze.used[pos.row + 1][pos.col] != 0:
+        possible_dirs.remove(Direction.DOWN)
+    return possible_dirs
 
 
 def get_maze_dimensions() -> Tuple[int, int]:
     while True:
-        width_str, length_str = input("What are your width and length?").split(",")
+        input_str = input("What are your width and length?")
-        width = int(width_str)
+        if input_str.count(",") == 1:
-        length = int(length_str)
+            width_str, length_str = input_str.split(",")
-        if width > 1 and length > 1:
+            width = int(width_str)
-            break
+            length = int(length_str)
+            if width > 1 and length > 1:
+                break
         print("Meaningless dimensions. Try again.")
     return width, length
 
 
-def print_maze(maze: Maze) -> None:
-    for col in range(maze.width):
-        if col == maze.enter_col:
-            print(".  ", end="")
-        else:
-            print(".--", end="")
-    print(".")
-    for row in range(maze.length):
-        print("I", end="")
-        for col in range(maze.width):
-            if maze.walls[row][col] < 2:
-                print("  I", end="")
-            else:
-                print("   ", end="")
-        print()
-        for col in range(maze.width):
-            if maze.walls[row][col] == 0 or maze.walls[row][col] == 2:
-                print(":--", end="")
-            else:
-                print(":  ", end="")
-        print(".")
-
-
 if __name__ == "__main__":
     main()

02_Amazing/python/test_amazing.py

@@ -1,8 +1,13 @@
+import io
 import pytest
-from amazing import build_maze, welcome_header
+from _pytest.capture import CaptureFixture
+from _pytest.monkeypatch import MonkeyPatch
+
+from amazing import build_maze, welcome_header, main
 
 
-def test_welcome_header(capsys) -> None:
+def test_welcome_header(capsys: CaptureFixture[str]) -> None:
+    capsys.readouterr()
     welcome_header()
     out, err = capsys.readouterr()
     assert out == (
@@ -25,3 +30,29 @@ def test_welcome_header(capsys) -> None:
 def test_build_maze(width: int, length: int) -> None:
     with pytest.raises(AssertionError):
         build_maze(width, length)
+
+
+@pytest.mark.parametrize(
+    ("width", "length"),
+    [
+        (3, 3),
+        (10, 10),
+    ],
+)
+def test_main(monkeypatch: MonkeyPatch, width: int, length: int) -> None:
+    monkeypatch.setattr(
+        "sys.stdin",
+        io.StringIO(f"{width},{length}"),
+    )
+    main()
+
+
+def test_main_error(monkeypatch: MonkeyPatch) -> None:
+    width = 1
+    length = 2
+
+    monkeypatch.setattr(
+        "sys.stdin",
+        io.StringIO(f"{width},{length}\n3,3"),
+    )
+    main()