add types to function args

snakes/BestBattleSnake.py

@@ -1,12 +1,17 @@
+from collections.abc import Iterator
 from collections import deque
 from typing import Any, cast
-import random
-import os
+import random, os
 
 from snakes.TemplateSnake import TemplateSnake
 
 class BestBattleSnake(TemplateSnake):
   VERSION = "2.5.0"
+  Point = tuple[int, int]
+  Coord = dict[str, int]
+  SnakeState = dict[str, Any]
+  MoveMap = dict[str, Coord]
+  AttackMap = dict[Point, int]
 
   DIRECTIONS = {
     "up": (0, 1),
@@ -32,7 +37,7 @@ class BestBattleSnake(TemplateSnake):
     self.previous_hazards = set()
     self.duel_style = self._get_duel_style()
 
-  def _get_duel_style(self):
+  def _get_duel_style(self) -> str:
     """Resolve duel tuning style from `BATTLE_SNAKE_DUEL_STYLE` or `DUEL_STYLE`."""
     value = os.getenv("BATTLE_SNAKE_DUEL_STYLE")
     if value is None:
@@ -43,7 +48,7 @@ class BestBattleSnake(TemplateSnake):
       return "balanced"
     return style
 
-  def _duel_weights(self, style):
+  def _duel_weights(self, style:str) -> dict[str, float]:
     """Return score multipliers for the selected duel style preset."""
     if style == "safe":
       return {
@@ -63,7 +68,7 @@ class BestBattleSnake(TemplateSnake):
       "food_bias": 1.00,
     }
 
-  def choose_move(self, game_data):
+  def choose_move(self, game_data:dict) -> str:
     """Pick the next move from a Battlesnake move request.
 
     Docs: https://docs.battlesnake.com/api/example-move
@@ -185,8 +190,8 @@ class BestBattleSnake(TemplateSnake):
       self.previous_hazards = set(hazard_set)
       return best_move
 
-    scores: dict[str, float] = {}
-    move_safety: dict[str, dict[str, Any]] = {}
+    scores:dict[str, float] = {}
+    move_safety:dict[str, dict[str, Any]] = {}
     for move, pos in safe_moves.items():
       point = (pos["x"], pos["y"])
       ate_food = point in food_set
@@ -336,7 +341,7 @@ class BestBattleSnake(TemplateSnake):
     self.previous_hazards = set(hazard_set)
     return best_move
 
-  def _choose_duel_move(self, safe_moves, my_body, my_len, my_health, food_set, hazard_set, other_snakes, enemy_attack_map, enemy_can_grow_cache, previous_hazard_set, hazard_damage, width, height):
+  def _choose_duel_move(self, safe_moves:MoveMap, my_body:list[Coord], my_len:int, my_health:int, food_set:set[Point], hazard_set:set[Point], other_snakes:list[SnakeState],enemy_attack_map:AttackMap, enemy_can_grow_cache:dict[Any, bool], previous_hazard_set:set[Point], hazard_damage:int, width:int, height:int) -> tuple[str, dict[str, float]]:
     """Score and select a move for one-vs-one games."""
     duel_weights = self._duel_weights(self.duel_style)
     enemy = other_snakes[0]
@@ -344,8 +349,8 @@ class BestBattleSnake(TemplateSnake):
     enemy_len = enemy.get("length", len(enemy["body"]))
     encase_target_space = max(8, enemy_len * 2)
 
-    scores: dict[str, float] = {}
-    move_safety: dict[str, dict[str, Any]] = {}
+    scores:dict[str, float] = {}
+    move_safety:dict[str, dict[str, Any]] = {}
 
     for move, pos in safe_moves.items():
       point = (pos["x"], pos["y"])
@@ -507,10 +512,10 @@ class BestBattleSnake(TemplateSnake):
     ]
     return random.choice(top_moves), scores
 
-  def _choose_constrictor_move(self, safe_moves, my_body, my_len, other_snakes, food_set, enemy_attack_map, enemy_heads, enemy_can_grow_cache, width, height):
+  def _choose_constrictor_move(self, safe_moves:MoveMap, my_body:list[Coord], my_len:int, other_snakes:list[SnakeState], food_set:set[Point], enemy_attack_map:AttackMap, enemy_heads:list[Point], enemy_can_grow_cache:dict[Any, bool], width:int, height:int) -> tuple[str, dict[str, float]]:
     """Score and select a move for constrictor games."""
-    scores: dict[str, float] = {}
-    move_safety: dict[str, dict[str, Any]] = {}
+    scores:dict[str, float] = {}
+    move_safety:dict[str, dict[str, Any]] = {}
 
     for move, pos in safe_moves.items():
       point = (pos["x"], pos["y"])
@@ -602,7 +607,7 @@ class BestBattleSnake(TemplateSnake):
     ]
     return random.choice(top_moves), scores
 
-  def _legal_moves(self, my_head, my_body, other_snakes, food_set, is_constrictor, width, height):
+  def _legal_moves(self, my_head:Coord, my_body:list[Coord], other_snakes:list[SnakeState], food_set:set[Point], is_constrictor:bool, width:int, height:int) -> MoveMap:
     """Return legal immediate moves after body, wall, and tail checks."""
     occupied = self._occupied_cells(my_body, other_snakes)
     own_tail = (my_body[-1]["x"], my_body[-1]["y"])
@@ -630,14 +635,14 @@ class BestBattleSnake(TemplateSnake):
 
     return safe_moves
 
-  def _occupied_cells(self, my_body, other_snakes):
+  def _occupied_cells(self, my_body:list[Coord], other_snakes:list[SnakeState]) -> set[Point]:
     """Build a set of occupied coordinates for all snake bodies."""
     occupied = {(segment["x"], segment["y"]) for segment in my_body}
     for snake in other_snakes:
       occupied |= {(segment["x"], segment["y"]) for segment in snake["body"]}
     return occupied
 
-  def _simulation_blocked(self, future_body, other_snakes, food_set, is_constrictor, enemy_can_grow_cache=None):
+  def _simulation_blocked(self, future_body:list[Coord], other_snakes:list[SnakeState], food_set:set[Point], is_constrictor:bool, enemy_can_grow_cache:dict[Any, bool]|None=None) -> set[Point]:
     """Build blocked cells for evaluating the board one turn ahead."""
     blocked = {(segment["x"], segment["y"]) for segment in future_body}
 
@@ -669,7 +674,7 @@ class BestBattleSnake(TemplateSnake):
 
     return blocked
 
-  def _build_enemy_attack_map(self, my_snake, other_snakes, food_set, is_constrictor, width, height, enemy_can_grow_cache=None):
+  def _build_enemy_attack_map(self, my_snake:SnakeState, other_snakes:list[SnakeState], food_set:set[Point], is_constrictor:bool, width:int, height:int, enemy_can_grow_cache:dict[Any, bool]|None=None) -> AttackMap:
     """Map cells enemies can contest next turn to their effective length."""
     occupied = self._occupied_cells(my_snake["body"], other_snakes)
     my_body_points = {(segment["x"], segment["y"]) for segment in my_snake["body"]}
@@ -712,7 +717,7 @@ class BestBattleSnake(TemplateSnake):
 
     return attack_map
 
-  def _future_body(self, current_body, next_head, ate_food, is_constrictor):
+  def _future_body(self, current_body:list[Coord], next_head:Coord, ate_food:bool, is_constrictor:bool) -> list[Coord]:
     """Simulate future body segments after a candidate move."""
     next_body = [next_head]
     next_body.extend(current_body)
@@ -723,7 +728,7 @@ class BestBattleSnake(TemplateSnake):
     next_body.pop()
     return next_body
 
-  def _can_step_on_own_tail(self, point, own_tail, own_tail_is_stacked, ate_food, is_constrictor):
+  def _can_step_on_own_tail(self, point:Point, own_tail:Point, own_tail_is_stacked:bool, ate_food:bool, is_constrictor:bool) -> bool:
     """Return whether stepping onto our tail is allowed this turn."""
     if is_constrictor:
       return False
@@ -733,13 +738,13 @@ class BestBattleSnake(TemplateSnake):
       return False
     return point == own_tail
 
-  def _is_tail_stacked(self, body):
+  def _is_tail_stacked(self, body:list[Coord]) -> bool:
     """Check whether tail overlaps the previous body segment."""
     if len(body) < 2:
       return False
     return body[-1]["x"] == body[-2]["x"] and body[-1]["y"] == body[-2]["y"]
 
-  def _enemy_can_grow_this_turn(self, snake, food_set):
+  def _enemy_can_grow_this_turn(self, snake:SnakeState, food_set:set[Point]) -> bool:
     """Return True if an enemy can eat food in one move."""
     head = snake["head"]
     for dx, dy in self.DIRECTIONS.values():
@@ -747,7 +752,7 @@ class BestBattleSnake(TemplateSnake):
         return True
     return False
 
-  def _hazard_damage_per_turn(self, game_data):
+  def _hazard_damage_per_turn(self, game_data:dict) -> int:
     """Read royale hazard damage from ruleset settings.
 
     Docs: https://docs.battlesnake.com/maps/royale
@@ -760,7 +765,7 @@ class BestBattleSnake(TemplateSnake):
     settings = ruleset.get("settings", {}) if isinstance(ruleset, dict) else {}
     return int(settings.get("hazardDamagePerTurn", 15))
 
-  def _hazard_is_active(self, point, ate_food, hazard_set, previous_hazard_set):
+  def _hazard_is_active(self, point:Point, ate_food:bool, hazard_set:set[Point], previous_hazard_set:set[Point]) -> bool:
     """Apply royale hazard grace and food-exception behavior.
 
     Docs: https://docs.battlesnake.com/maps/royale
@@ -771,7 +776,7 @@ class BestBattleSnake(TemplateSnake):
       return False
     return point in previous_hazard_set
 
-  def _nearest_food_distance(self, start, food_set, blocked, width, height):
+  def _nearest_food_distance(self, start:Point, food_set:set[Point], blocked:set[Point], width:int, height:int) -> int|None:
     """Compute shortest reachable distance to any food using BFS."""
     if not food_set:
       return None
@@ -795,7 +800,7 @@ class BestBattleSnake(TemplateSnake):
 
     return None
 
-  def _path_distance(self, start, goal, blocked, width, height):
+  def _path_distance( self, start:Point, goal:Point, blocked:set[Point], width:int, height:int) -> int|None:
     """Compute shortest path distance between two cells."""
     queue = deque([(start, 0)])
     seen = {start}
@@ -817,7 +822,7 @@ class BestBattleSnake(TemplateSnake):
 
     return None
 
-  def _flood_fill_count(self, start, blocked, width, height):
+  def _flood_fill_count(self, start:Point, blocked:set[Point], width:int, height:int) -> int:
     """Count reachable cells from `start` using flood fill."""
     queue = deque([start])
     seen = {start}
@@ -836,7 +841,7 @@ class BestBattleSnake(TemplateSnake):
 
     return len(seen)
 
-  def _open_neighbor_count(self, start, blocked, width, height):
+  def _open_neighbor_count(self, start:Point, blocked:set[Point], width:int, height:int) -> int:
     """Count walkable orthogonal neighbors around `start`."""
     count = 0
     for neighbor in self._neighbors(start):
@@ -847,7 +852,7 @@ class BestBattleSnake(TemplateSnake):
       count += 1
     return count
 
-  def _next_turn_option_count(self, future_body, blocked, width, height):
+  def _next_turn_option_count(self, future_body:list[Coord], blocked:set[Point], width:int, height:int) -> int:
     """Estimate options available after the next simulated turn."""
     if not future_body:
       return 0
@@ -863,7 +868,7 @@ class BestBattleSnake(TemplateSnake):
       count += 1
     return count
 
-  def _revisit_penalty(self, point):
+  def _revisit_penalty(self, point:Point) -> float:
     """Return penalty for revisiting recent head positions."""
     if not self.recent_heads:
       return 0.0
@@ -874,7 +879,7 @@ class BestBattleSnake(TemplateSnake):
       penalty += max(0.0, 18.0 - index * 2.0)
     return penalty
 
-  def _territory_control_score(self, my_start, enemy_starts, blocked, width, height):
+  def _territory_control_score(self, my_start:Point, enemy_starts:list[Point], blocked:set[Point], width:int, height:int) -> int:
     """Estimate territorial advantage versus enemy start positions."""
     if not enemy_starts:
       return 0
@@ -910,10 +915,10 @@ class BestBattleSnake(TemplateSnake):
 
     return score
 
-  def _distance_map(self, start, blocked, width, height):
+  def _distance_map(self, start:Point, blocked:set[Point], width:int, height:int) -> dict[Point, int]:
     """Build a BFS distance map from the given start cell."""
     queue = deque([(start, 0)])
-    distances = {start: 0}
+    distances = {start:0}
 
     while queue:
       point, distance = queue.popleft()
@@ -929,7 +934,7 @@ class BestBattleSnake(TemplateSnake):
 
     return distances
 
-  def _enemy_confinement_metrics(self, enemy_head, blocked, width, height):
+  def _enemy_confinement_metrics(self, enemy_head:Point, blocked:set[Point], width:int, height:int) -> tuple[int, int]:
     """Return enemy reachable space and immediate exit count."""
     enemy_blocked = set(blocked)
     enemy_blocked.discard(enemy_head)
@@ -939,20 +944,20 @@ class BestBattleSnake(TemplateSnake):
     )
     return enemy_space, enemy_options
 
-  def _neighbors(self, point):
+  def _neighbors(self, point:Point) -> Iterator[Point]:
     """Yield orthogonal neighbor coordinates for a point."""
     for dx, dy in self.DIRECTIONS.values():
       yield (point[0] + dx, point[1] + dy)
 
-  def _manhattan(self, a, b):
+  def _manhattan(self, a:Point, b:Point) -> int:
     """Return Manhattan distance between two points."""
     return abs(a[0] - b[0]) + abs(a[1] - b[1])
 
-  def _in_bounds(self, point, width, height):
+  def _in_bounds(self, point:Point, width:int, height:int) -> bool:
     """Return True when a point is inside board boundaries."""
     return 0 <= point[0] < width and 0 <= point[1] < height
 
-  def _fallback_move(self, head, width, height):
+  def _fallback_move(self, head:Coord, width:int, height:int) -> str:
     """Pick the first in-bounds move as emergency fallback."""
     for move, (dx, dy) in self.DIRECTIONS.items():
       point = (head["x"] + dx, head["y"] + dy)