diff --git a/socha/api/plugin/penguins.py b/socha/api/plugin/penguins.py index 4a7db62..3ae3324 100644 --- a/socha/api/plugin/penguins.py +++ b/socha/api/plugin/penguins.py @@ -676,7 +676,7 @@ def contains_all(self, fields: List[Field]) -> bool: return False return True - def get_moves_in_direction(self, origin: HexCoordinate, direction: Vector, team_enum: TeamEnum) -> List[Move]: + def get_moves_in_direction(self, origin: HexCoordinate, direction: Vector, team_enum: Optional[TeamEnum] = None) -> List[Move]: """ Gets all moves in the given direction from the given origin. @@ -689,6 +689,8 @@ def get_moves_in_direction(self, origin: HexCoordinate, direction: Vector, team_ List[Move]: List of moves that can be made in the given direction from the given index, for the given team_enum """ + if team_enum is None: + team_enum = self.get_field(origin).penguin.team_enum if not self.get_field(origin).penguin or self.get_field(origin).penguin.team_enum != team_enum: return [] @@ -705,7 +707,7 @@ def _is_destination_valid(self, field: HexCoordinate) -> bool: return self.is_valid(field) and not self.is_occupied(field) and not \ self.get_field(field).is_empty() - def possible_moves_from(self, position: HexCoordinate, team_enum: TeamEnum) -> List[Move]: + def possible_moves_from(self, position: HexCoordinate, team_enum: Optional[TeamEnum] = None) -> List[Move]: """ Returns a list of all possible moves from the given position. That are all moves in all hexagonal directions. @@ -721,7 +723,7 @@ def possible_moves_from(self, position: HexCoordinate, team_enum: TeamEnum) -> L """ if not self.is_valid(position): raise IndexError(f"Index out of range: [x={position.x}, y={position.y}]") - if not self.get_field(position).penguin or self.get_field(position).penguin.team_enum != team_enum: + if not self.get_field(position).penguin or (team_enum != None and self.get_field(position).penguin.team_enum != team_enum): return [] return [move for direction in Vector().directions for move in self.get_moves_in_direction(position, direction, team_enum)]