WIP: Spec the prepare_bot dicts

Author: Debilski

pelita/game.py

@@ -16,7 +16,7 @@
 from .gamestate_filters import noiser, relocate_expired_food, update_food_age, in_homezone
 from .layout import get_legal_positions, initial_positions
 from .network import Controller, RemotePlayerFailure, RemotePlayerRecvTimeout, RemotePlayerSendError, ZMQPublisher
-from .spec import GameState, Layout, Pos
+from .spec import GameState, Layout, Pos, TeamInitial, TeamState, TeamStateFinished
 from .team import RemoteTeam, make_team
 from .viewer import (AsciiViewer, ProgressViewer, ReplayWriter, ReplyToViewer,
                      ResultPrinter)
@@ -660,89 +660,98 @@ def request_new_position(game_state):
     return bot_reply
 
 
-def prepare_bot_state(game_state, team_idx=None):
+def prepare_bot_state(game_state: GameState, team_idx=None) -> TeamState | TeamInitial | TeamStateFinished:
     """ Prepares the bot’s game state for the current bot.
 
     NB: This will update the game_state to store new noisy positions.
     """
-    if game_state['game_phase'] == 'INIT':
-        # We assume that we are in get_initial phase
-        turn = team_idx
-        bot_turn = None
-        seed = game_state['rng'].randint(0, sys.maxsize)
-    elif game_state['game_phase'] == 'FINISHED':
-        # Called for remote players in _exit
-        turn = team_idx
-        bot_turn = None
-        seed = None
-    elif game_state['game_phase'] == 'RUNNING':
-        turn = game_state['turn']
-        bot_turn = game_state['turn'] // 2
-        seed = None
-    else:
-        _logger.warning("Got bad game_state in prepare_bot_state")
-        return
-
-    bot_position = game_state['bots'][turn]
-    own_team = turn % 2
-    enemy_team = 1 - own_team
-    enemy_positions = game_state['bots'][enemy_team::2]
-    noised_positions = noiser(walls=game_state['walls'],
-                              shape=game_state['shape'],
-                              bot_position=bot_position,
-                              enemy_positions=enemy_positions,
-                              noise_radius=game_state['noise_radius'],
-                              sight_distance=game_state['sight_distance'],
-                              rng=game_state['rng'])
-
-
-    # Update noisy_positions in the game_state
-    # reset positions
-    game_state['noisy_positions'] = [None] * 4
-    noisy_or_none = [
-        noisy_pos if is_noisy else None
-            for is_noisy, noisy_pos in
-            zip(noised_positions['is_noisy'], noised_positions['enemy_positions'])
-    ]
-    game_state['noisy_positions'][enemy_team::2] = noisy_or_none
-
-    bots = game_state['bots'][:]
-    bots[enemy_team::2] = noised_positions['enemy_positions']
-
-    is_noisy = [False for _ in range(4)]
-    is_noisy[enemy_team::2] = noised_positions['is_noisy']
-
-    shaded_food_own = list(pos for pos, age in game_state['food_age'][own_team].items()
-                       if age > 0)
-    shaded_food = [[], []]
-    shaded_food[own_team] = shaded_food_own
-
-    bot_state = {
-        'bots': bots,
-        'score': game_state['score'][:],
-        'kills': game_state['kills'][:],
-        'deaths': game_state['deaths'][:],
-        'bot_was_killed': game_state['bot_was_killed'][:],
-        'error_count': [len(e) for e in game_state['timeouts'][:]],
-        'food': [list(team_food) for team_food in game_state['food']],
-        'shaded_food': shaded_food,
-        'team_time': game_state['team_time'][:],
-        'is_noisy': is_noisy,
-        'round': game_state['round'],
-        'turn': turn,
-        'timeout_length': game_state['timeout_length'],
-    }
+    match game_state['game_phase']:
+        case "INIT":
+            turn = team_idx
+            seed = game_state['rng'].randint(0, sys.maxsize)
+
+            team_state_initial: TeamInitial = {
+                'walls': game_state['walls'],
+                'shape': game_state['shape'],
+                'seed': seed,
+                'max_rounds': game_state['max_rounds'],
+                'team_names': game_state['team_names'][:],
+                'timeout_length': game_state['timeout_length'],
+            }
+            return team_state_initial
+        case "RUNNING":
+
+            turn = game_state['turn']
+
+            bot_position = game_state['bots'][turn]
+            own_team = turn % 2
+            enemy_team = 1 - own_team
+            enemy_positions = game_state['bots'][enemy_team::2]
+            noised_positions = noiser(walls=game_state['walls'],
+                                    shape=game_state['shape'],
+                                    bot_position=bot_position,
+                                    enemy_positions=enemy_positions,
+                                    noise_radius=game_state['noise_radius'],
+                                    sight_distance=game_state['sight_distance'],
+                                    rng=game_state['rng'])
+
+            # Update noisy_positions in the game_state
+            # reset positions
+            game_state['noisy_positions'] = [None] * 4
+            noisy_or_none = [
+                noisy_pos if is_noisy else None
+                    for is_noisy, noisy_pos in
+                    zip(noised_positions['is_noisy'], noised_positions['enemy_positions'])
+            ]
+            game_state['noisy_positions'][enemy_team::2] = noisy_or_none
+
+            bots = game_state['bots'][:]
+            bots[enemy_team::2] = noised_positions['enemy_positions']
+
+            is_noisy = [False for _ in range(4)]
+            is_noisy[enemy_team::2] = noised_positions['is_noisy']
+
+            shaded_food_own = list(pos for pos, age in game_state['food_age'][own_team].items()
+                            if age > 0)
+            shaded_food = [[], []]
+            shaded_food[own_team] = shaded_food_own
+
+            bot_state: TeamState = {
+                'bots': bots,
+                'score': game_state['score'][:],
+                'kills': game_state['kills'][:],
+                'deaths': game_state['deaths'][:],
+                'bot_was_killed': game_state['bot_was_killed'][:],
+                'error_count': [len(e) for e in game_state['timeouts'][:]],
+                'food': [list(team_food) for team_food in game_state['food']],
+                'shaded_food': shaded_food,
+                'team_time': game_state['team_time'][:],
+                'is_noisy': is_noisy,
+                'round': game_state['round'],
+                'turn': turn,
+            }
+            return bot_state
+
+        case "FINISHED":
+            # Called for remote players in _exit
+            turn = team_idx
+            # TODO: is turn needed?
+
+            team_state_final: TeamStateFinished = {
+                'bots': game_state['bots'][:],
+                'score': game_state['score'][:],
+                'kills': game_state['kills'][:],
+                'deaths': game_state['deaths'][:],
+                'bot_was_killed': game_state['bot_was_killed'][:],
+                'error_count': [len(e) for e in game_state['timeouts'][:]],
+                'food': [list(team_food) for team_food in game_state['food']],
+                'team_time': game_state['team_time'][:],
+                'round': game_state['round'],
+                'turn': turn,
+            }
+            return team_state_final
 
-    if game_state['game_phase'] == 'INIT':
-        bot_state.update({
-            'walls': game_state['walls'], # only in initial round
-            'shape': game_state['shape'], # only in initial round
-            'seed': seed, # only used in set_initial phase
-            'max_rounds': game_state['max_rounds'],
-            'team_names': game_state['team_names'][:],
-        })
-
-    return bot_state
+    raise PelitaIllegalGameState("Got bad game_state in prepare_bot_state")
 
 
 def update_viewers(game_state):

pelita/spec.py

@@ -49,3 +49,38 @@ class GameState(TypedDict):
     error_limit: int
     viewers: list[Any]
     controller: None|Any
+
+class TeamInitial(TypedDict):
+    walls: set[Pos]
+    shape: Shape
+    seed: int
+    max_rounds: int
+    team_names: list[str]
+    timeout_length: float
+
+class TeamState(TypedDict):
+    bots: list[Pos]
+    score: list[int]
+    kills: list[int]
+    deaths: list[int]
+    bot_was_killed: list[bool]
+    error_count: list[int]
+    food: list[list[Pos]]
+    shaded_food: list[list[Pos]]
+    team_time: list[float]
+    is_noisy: list[bool]
+    round: int|None
+    turn: int
+
+class TeamStateFinished(TypedDict):
+    bots: list[Pos]
+    score: list[int]
+    kills: list[int]
+    deaths: list[int]
+    bot_was_killed: list[bool]
+    error_count: list[bool]
+    food: list[list[Pos]]
+    team_time: list[float]
+    round: int|None
+    turn: int
+

pelita/team.py

@@ -391,7 +391,9 @@ def wait_ready(self, timeout):
             raise RemotePlayerRecvTimeout("", "") from None
 
     def set_initial(self, team_id, game_state):
+        # TODO: timeout length should be set when object is created
         timeout_length = game_state['timeout_length']
+        self.request_timeout = timeout_length
 
         msg_id = self.conn.send_req("set_initial", {"team_id": team_id,
                                                 "game_state": game_state})
@@ -401,7 +403,7 @@ def set_initial(self, team_id, game_state):
         return reply
 
     def get_move(self, game_state):
-        timeout_length = game_state['timeout_length']
+        timeout_length = self.request_timeout
 
         msg_id = self.conn.send_req("get_move", {"game_state": game_state})
         reply = self.conn.recv_reply(msg_id, timeout_length)