Python: Add tests and type annotations

08_Batnum/python/batnum.py

@@ -1,18 +1,44 @@
-from enum import Enum
-from typing import Tuple, Union
+from enum import IntEnum
+from typing import Tuple, Any
 
 
-class WinOptions(Enum):
+class WinOptions(IntEnum):
     Undefined = 0
     TakeLast = 1
     AvoidLast = 2
 
+    @classmethod
+    def _missing_(cls, value: Any) -> "WinOptions":
+        try:
+            int_value = int(value)
+        except Exception:
+            return WinOptions.Undefined
+        if int_value == 1:
+            return WinOptions.TakeLast
+        elif int_value == 2:
+            return WinOptions.AvoidLast
+        else:
+            return WinOptions.Undefined
 
-class StartOptions(Enum):
+
+class StartOptions(IntEnum):
     Undefined = 0
     ComputerFirst = 1
     PlayerFirst = 2
 
+    @classmethod
+    def _missing_(cls, value: Any) -> "StartOptions":
+        try:
+            int_value = int(value)
+        except Exception:
+            return StartOptions.Undefined
+        if int_value == 1:
+            return StartOptions.ComputerFirst
+        elif int_value == 2:
+            return StartOptions.PlayerFirst
+        else:
+            return StartOptions.Undefined
+
 
 def print_intro() -> None:
     """Prints out the introduction and rules for the game."""
@@ -34,7 +60,7 @@ def print_intro() -> None:
     return
 
 
-def get_params() -> Tuple[int, int, int, int, int]:
+def get_params() -> Tuple[int, int, int, StartOptions, WinOptions]:
     """This requests the necessary parameters to play the game.
 
     Returns a set with the five game parameters:
@@ -46,46 +72,69 @@ def get_params() -> Tuple[int, int, int, int, int]:
         winOption - 1 if the goal is to take the last object
                     or 2 if the goal is to not take the last object
     """
+    pile_size = get_pile_size()
+    if pile_size < 0:
+        return (-1, 0, 0, StartOptions.Undefined, WinOptions.Undefined)
+    win_option = get_win_option()
+    min_select, max_select = get_min_max()
+    start_option = get_start_option()
+    return (pile_size, min_select, max_select, start_option, win_option)
+
+
+def get_pile_size() -> int:
+    # A negative number will stop the game.
     pile_size = 0
-    win_option: Union[WinOptions, int] = WinOptions.Undefined
+    while pile_size == 0:
+        try:
+            pile_size = int(input("ENTER PILE SIZE "))
+        except ValueError:
+            pile_size = 0
+    return pile_size
+
+
+def get_win_option() -> WinOptions:
+    win_option: WinOptions = WinOptions.Undefined
+    while win_option == WinOptions.Undefined:
+        win_option = WinOptions(input("ENTER WIN OPTION - 1 TO TAKE LAST, 2 TO AVOID LAST: "))  # type: ignore
+    return win_option
+
+
+def get_min_max() -> Tuple[int, int]:
     min_select = 0
     max_select = 0
-    start_option: Union[StartOptions, int] = StartOptions.Undefined
-
-    while pile_size < 1:
-        pile_size = int(input("ENTER PILE SIZE "))
-    while win_option == WinOptions.Undefined:
-        win_option = int(input("ENTER WIN OPTION - 1 TO TAKE LAST, 2 TO AVOID LAST: "))
-    assert isinstance(win_option, int)
     while min_select < 1 or max_select < 1 or min_select > max_select:
         (min_select, max_select) = (
             int(x) for x in input("ENTER MIN AND MAX ").split(" ")
         )
+    return min_select, max_select
+
+
+def get_start_option() -> StartOptions:
+    start_option: StartOptions = StartOptions.Undefined
     while start_option == StartOptions.Undefined:
-        start_option = int(input("ENTER START OPTION - 1 COMPUTER FIRST, 2 YOU FIRST "))
-    assert isinstance(start_option, int)
-    return (pile_size, min_select, max_select, start_option, win_option)
+        start_option = StartOptions(input("ENTER START OPTION - 1 COMPUTER FIRST, 2 YOU FIRST "))  # type: ignore
+    return start_option
 
 
 def player_move(
-    pile_size, min_select, max_select, start_option, win_option
+    pile_size: int, min_select: int, max_select: int, win_option: WinOptions
 ) -> Tuple[bool, int]:
     """This handles the player's turn - asking the player how many objects
     to take and doing some basic validation around that input.  Then it
     checks for any win conditions.
 
     Returns a boolean indicating whether the game is over and the new pileSize."""
-    playerDone = False
-    while not playerDone:
-        playerMove = int(input("YOUR MOVE "))
-        if playerMove == 0:
+    player_done = False
+    while not player_done:
+        player_move = int(input("YOUR MOVE "))
+        if player_move == 0:
             print("I TOLD YOU NOT TO USE ZERO!  COMPUTER WINS BY FORFEIT.")
             return (True, pile_size)
-        if playerMove > max_select or playerMove < min_select:
+        if player_move > max_select or player_move < min_select:
             print("ILLEGAL MOVE, REENTER IT")
             continue
-        pile_size = pile_size - playerMove
-        playerDone = True
+        pile_size = pile_size - player_move
+        player_done = True
         if pile_size <= 0:
             if win_option == WinOptions.AvoidLast:
                 print("TOUGH LUCK, YOU LOSE.")
@@ -95,7 +144,9 @@ def player_move(
     return (False, pile_size)
 
 
-def computer_pick(pile_size, min_select, max_select, start_option, win_option) -> int:
+def computer_pick(
+    pile_size: int, min_select: int, max_select: int, win_option: WinOptions
+) -> int:
     """This handles the logic to determine how many objects the computer
     will select on its turn.
     """
@@ -110,7 +161,7 @@ def computer_pick(pile_size, min_select, max_select, start_option, win_option) -
 
 
 def computer_move(
-    pile_size, min_select, max_select, start_option, win_option
+    pile_size: int, min_select: int, max_select: int, win_option: WinOptions
 ) -> Tuple[bool, int]:
     """This handles the computer's turn - first checking for the various
     win/lose conditions and then calculating how many objects
@@ -132,13 +183,19 @@ def computer_move(
         return (True, pile_size)
 
     # Otherwise, we determine how many the computer selects
-    currSel = computer_pick(pile_size, min_select, max_select, start_option, win_option)
-    pile_size = pile_size - currSel
-    print(f"COMPUTER TAKES {currSel} AND LEAVES {pile_size}")
+    curr_sel = computer_pick(pile_size, min_select, max_select, win_option)
+    pile_size = pile_size - curr_sel
+    print(f"COMPUTER TAKES {curr_sel} AND LEAVES {pile_size}")
     return (False, pile_size)
 
 
-def play_game(pile_size, min_select, max_select, start_option, win_option) -> None:
+def play_game(
+    pile_size: int,
+    min_select: int,
+    max_select: int,
+    start_option: StartOptions,
+    win_option: WinOptions,
+) -> None:
     """This is the main game loop - repeating each turn until one
     of the win/lose conditions is met.
     """
@@ -150,30 +207,29 @@ def play_game(pile_size, min_select, max_select, start_option, win_option) -> No
     while not game_over:
         if players_turn:
             (game_over, pile_size) = player_move(
-                pile_size, min_select, max_select, start_option, win_option
+                pile_size, min_select, max_select, win_option
             )
             players_turn = False
             if game_over:
                 return
         if not players_turn:
             (game_over, pile_size) = computer_move(
-                pile_size, min_select, max_select, start_option, win_option
+                pile_size, min_select, max_select, win_option
             )
             players_turn = True
 
 
-if __name__ == "__main__":
-
-    pileSize = 0
-    minSelect = 0
-    maxSelect = 0
-    # 1 = to take last, 2 = to avoid last
-    winOption = 0
-    # 1 = computer first, 2 = user first
-    startOption = 0
-
+def main() -> None:
     while True:
         print_intro()
-        (pileSize, minSelect, maxSelect, startOption, winOption) = get_params()
+        (pile_size, min_select, max_select, start_option, win_option) = get_params()
+
+        if pile_size < 0:
+            return
+
         # Just keep playing the game until the user kills it with ctrl-C
-        play_game(pileSize, minSelect, maxSelect, startOption, winOption)
+        play_game(pile_size, min_select, max_select, start_option, win_option)
+
+
+if __name__ == "__main__":
+    main()

08_Batnum/python/test_batnum.py

@@ -0,0 +1,21 @@
+import io
+from _pytest.capture import CaptureFixture
+from _pytest.monkeypatch import MonkeyPatch
+
+from batnum import main
+
+
+def test_main_win(monkeypatch: MonkeyPatch, capsys: CaptureFixture[str]) -> None:
+    pile_size = 1
+    monkeypatch.setattr("sys.stdin", io.StringIO(f"{pile_size}\n1\n1 2\n2\n1\n-1\n"))
+    main()
+    captured = capsys.readouterr()
+    assert "CONGRATULATIONS, YOU WIN" in captured.out
+
+
+def test_main_lose(monkeypatch: MonkeyPatch, capsys: CaptureFixture[str]) -> None:
+    pile_size = 3
+    monkeypatch.setattr("sys.stdin", io.StringIO(f"{pile_size}\n2\n1 2\n2\n1\n1\n-1\n"))
+    main()
+    captured = capsys.readouterr()
+    assert "TOUGH LUCK, YOU LOSE" in captured.out