from enum import Enum, auto from dataclasses import dataclass from typing import Iterator, Optional, List, Tuple import random class PieceType(Enum): PAWN = auto() KNIGHT = auto() BISHOP = auto() ROOK = auto() QUEEN = auto() KING = auto() class Color(Enum): WHITE = auto() BLACK = auto() @property def opposite(self) -> "Color": if self == Color.WHITE: return Color.BLACK else: return Color.WHITE @dataclass class BoardPos: p: Tuple[int, int] def __str__(self) -> str: return f"({self.p[0]},{self.p[1]})" @property def x (self) -> int: return self.p[0] @property def y (self) -> int: return self.p[1] @dataclass class Piece: type: PieceType color: Color # uppercase = White, lowercase = Black def char(self) -> str: mapping = { PieceType.PAWN: "p", PieceType.KNIGHT: "n", PieceType.BISHOP: "b", PieceType.ROOK: "r", PieceType.QUEEN: "q", PieceType.KING: "k", } ch = mapping[self.type] return ch.upper() if self.color == Color.WHITE else ch def __str__(self) -> str: return self.char() def __repr__(self) -> str: return f"Piece({self.type.name},{self.color.name})" @dataclass class BoardMove: m_from: BoardPos m_to: BoardPos def __str__(self) -> str: return f"{self.m_from}->{self.m_to}" @dataclass class BoardField: piece: Optional[Piece] = None def __str__(self) -> str: return str(self.piece) if self.piece is not None else "." @dataclass class ChessBoard: fields: List[List[BoardField]] num_moves: int move_history: List[Tuple['BoardMove', Optional[Piece]]] def place(self, row: int, col: int, piece_type: PieceType, color: Color): self.fields[row][col].piece = Piece(piece_type, color) # initialize default starting position @classmethod def init_default(cls) -> "ChessBoard": brd = cls( [[BoardField() for _ in range(8)] for _ in range(8)], 0, [] ) # place pawns for col in range(8): brd.place(1, col, PieceType.PAWN, Color.BLACK) brd.place(6, col, PieceType.PAWN, Color.WHITE) # back rank order back_rank = [ PieceType.ROOK, PieceType.KNIGHT, PieceType.BISHOP, PieceType.QUEEN, PieceType.KING, PieceType.BISHOP, PieceType.KNIGHT, PieceType.ROOK, ] # place black back rank for col, piece_type in enumerate(back_rank): brd.place(0, col, piece_type, Color.BLACK) # place white back rank for col, piece_type in enumerate(back_rank): brd.place(7, col, piece_type, Color.WHITE) return brd # attempt to make move for color # returns true on success def make_move(self, move: BoardMove, color: Color) -> bool: sr, sc = move.m_from.p tr, tc = move.m_to.p src_field = self.fields[sr][sc] # validate move using legal moves for the player color legal = self.moves_basic(color) if not any(m.m_from.p == move.m_from.p and m.m_to.p == move.m_to.p for m in legal): return False # perform move and record history for unmaking captured = self.fields[tr][tc].piece self.fields[tr][tc].piece = src_field.piece src_field.piece = None self.move_history.append((move, captured)) self.num_moves += 1 return True # returns false if there is no move to unmake def unmake_move(self) -> bool: if not self.move_history or len(self.move_history) == 0: return False move, captured = self.move_history.pop() sr, sc = move.m_from.p tr, tc = move.m_to.p moving_piece = self.fields[tr][tc].piece # restore self.fields[sr][sc].piece = moving_piece self.fields[tr][tc].piece = captured self.num_moves = max(0, self.num_moves - 1) return True def moves_unchecked(self, piece: Piece, pos: BoardPos) -> List[List[Tuple[int, int]]]: row, column = pos.p rays: List[List[Tuple[int, int]]] = [] if piece.type == PieceType.KING: king_offsets = [(1, 1), (1, 0), (0, 1), (-1, -1), (-1, 0), (0, -1), (1, -1), (-1, 1)] rays.append([(row + dr, column + dc) for dr, dc in king_offsets]) elif piece.type == PieceType.KNIGHT: knight_offsets = [ (2, 1), (2, -1), (-2, 1), (-2, -1), (1, 2), (1, -2), (-1, 2), (-1, -2) ] rays.append([(row + dr, column + dc) for dr, dc in knight_offsets]) elif piece.type == PieceType.PAWN: # correct forward direction and start row depending on color if piece.color == Color.WHITE: forward = -1 start_row = 6 else: forward = 1 start_row = 1 forward_ray: List[Tuple[int, int]] = [] forward_ray.append((row + forward, column)) # if the pawn is at the start row, include the double-step if row == start_row: forward_ray.append((row + 2 * forward, column)) rays.append(forward_ray) # capture moves captures = [(row + forward, column - 1), (row + forward, column + 1)] rays.extend([[(rr, cc)] for rr, cc in captures]) elif piece.type == PieceType.BISHOP: directions = [(1, 1), (1, -1), (-1, 1), (-1, -1)] for dr, dc in directions: ray = [] for k in range(1, 8): ray.append((row + k * dr, column + k * dc)) rays.append(ray) elif piece.type == PieceType.ROOK: directions = [(1, 0), (-1, 0), (0, 1), (0, -1)] for dr, dc in directions: ray = [] for k in range(1, 8): ray.append((row + k * dr, column + k * dc)) rays.append(ray) elif piece.type == PieceType.QUEEN: directions = [ (1, 0), (-1, 0), (0, 1), (0, -1), (1, 1), (1, -1), (-1, 1), (-1, -1) ] for dr, dc in directions: ray = [] for k in range(1, 8): ray.append((row + k * dr, column + k * dc)) rays.append(ray) return rays def on_board(self, rr: int, cc: int) -> bool: return 0 <= rr < 8 and 0 <= cc < 8 # takes the color of the player whos possible moves will be returned def moves_basic(self, color: Color) -> List[BoardMove]: moves: List[BoardMove] = [] for pos, piece in self.iter_pieces(): if piece.color != color: continue rays = self.moves_unchecked(piece, pos) for ray in rays: # detect pawn capture ray and dont allow movement in that ray without capturing is_pawn_capture_ray = ( piece.type == PieceType.PAWN and len(ray) == 1 and ray[0][1] != pos.y and abs(ray[0][0] - pos.x) == 1 ) # non sliding rays get treated individually non_sliding = piece.type in (PieceType.KNIGHT, PieceType.KING) or is_pawn_capture_ray for (tr, tc) in ray: if not self.on_board(tr, tc): if non_sliding: continue break target_field = self.fields[tr][tc] if target_field.piece is None: # empty target if is_pawn_capture_ray: # pawn cant move diagonally into empty square break moves.append(BoardMove(pos, BoardPos((tr, tc)))) continue else: # occupied if target_field.piece.color == piece.color: break else: # opponent piece moves.append(BoardMove(pos, BoardPos((tr, tc)))) break # manually adding castling row_height = 0 if color == Color.BLACK: row_height = 7 expected_king_pos = BoardPos((row_height,4)) # if the king has moved, return as usual if self.has_piece_moved(expected_king_pos): return moves expected_lhs_rook_pos = BoardPos((row_height,0)) expected_rhs_rook_pos = BoardPos((row_height,7)) lhs_pieces_inbetween = [BoardPos((row_height, i)) for i in range(1, 3)] rhs_pieces_inbetween = [BoardPos((row_height, i)) for i in range(5, 6)] # make sure the pieces inbetween are gone and the rook has not moved if not self.has_piece_moved(expected_lhs_rook_pos) and self.are_pieces_none(lhs_pieces_inbetween): #todo: allow castling here if not self.has_piece_moved(expected_rhs_rook_pos) and self.are_pieces_none(rhs_pieces_inbetween): #todo: allow castling here return moves def has_piece_moved(self, pos: BoardPos) -> bool: for move in self.move_history: if move[0].m_from == pos: return True elif move[0].m_to == pos: return True return False def are_pieces_none(self, positions: List[BoardPos]): are_gone = True for pos in positions: if not self.is_piece_none(pos): are_gone = False return are_gone def is_piece_none(self, pos: BoardPos) -> bool: piece = self.fields[pos.x][pos.y].piece if piece is None: return True else: return False def is_piece_of_type(self, pos: BoardPos, piece_type: PieceType) -> bool: piece = self.fields[pos.x][pos.y].piece if piece is None: return False return piece.type == piece_type def iter_pieces(self) -> Iterator[Tuple[BoardPos, Piece]]: for row_idx, row in enumerate(self.fields): for col_idx, field in enumerate(row): if field.piece is not None: yield BoardPos((row_idx, col_idx)), field.piece def pos_of_king(self, color: Color) -> BoardPos: for pos, piece in self.iter_pieces(): if piece.type == PieceType.KING and piece.color == color: return pos raise ValueError("player has no king") # if a player does not have any moves, it has lost def generate_moves(self, color: Color) -> List[BoardMove]: # only moves after which the king is not in check are allowed # todo: the current method of checking what move will resolve the check is based on bruteforcing # maybe there is a better way of doing it, but for now this should suffice us_moves_wo_check = [] for move in self.moves_basic(color): # do a move and then check, if the king is still in check # if it isnt, add the move to the possibles ones if not self.make_move(move, color): raise ValueError("self moves basic created a move, which cannot be done (hopefully unreachable)") all_basic_enemy_moves = self.moves_basic(color.opposite) king_pos = self.pos_of_king(color) king_in_check = False for mv in all_basic_enemy_moves: if mv.m_to == king_pos: king_in_check = True if not king_in_check: us_moves_wo_check.append(move) if not self.unmake_move(): raise ValueError("failed to unmake move, shouldnt be an issue here") # check if our move lands directly on the enemies king, if it does, its illigal us_moves_rule_compliant = [] king_pos = self.pos_of_king(color.opposite) for move in us_moves_wo_check: if move.m_to != king_pos: us_moves_rule_compliant.append(move) return us_moves_wo_check def __str__(self) -> str: """ example string repr of starting pos: rnbqkbnr pppppppp ........ ........ ........ ........ PPPPPPPP RNBQKBNR """ lines: List[str] = [] for r in range(8): row_chars = [str(self.fields[r][c]) for c in range(8)] lines.append("".join(row_chars)) return "\n".join(lines) def play_random_game(board: Optional[ChessBoard] = None, max_moves: int = 400, verbose: bool = False) -> Tuple[ ChessBoard, int]: if board is None: board = ChessBoard.init_default() current = Color.WHITE # white moves first moves_played = 0 while moves_played < max_moves: legal = board.generate_moves(current) if not legal: # no legal moves if verbose: print(f"no legal moves for {current.name} after {moves_played} moves") return board, moves_played move = random.choice(legal) ok = board.make_move(move, current) # debug assert if not ok: # if an unexpected failure happens, break and treat as draw if verbose: print(f"make_move returned False for move {move} by {current}") return board, moves_played moves_played += 1 if verbose: print(f"{moves_played:03d}: {current.name} played {move}\nboard:\n{board}") # check whether the captured piece was a king _, captured = board.move_history[-1] if captured is not None and captured.type == PieceType.KING: if verbose: print(f"king captured by {current.name} on move {moves_played}.") raise RuntimeWarning("this should not be possible") return board, moves_played current = current.opposite if verbose: print(f"reached move limit ({max_moves})") return board, moves_played def run_random_games(n: int = 100, max_moves: int = 400, verbose: bool = False): for i in range(n): final_board, moves = play_random_game(None, max_moves=max_moves, verbose=verbose and (i < 3)) if verbose and (i < 3): print(f"game {i+1}: moves={moves}\nfinal_board:\n{final_board}") return # used only for testing purposesa def main(): default_brd = ChessBoard.init_default() default_brd.generate_moves # standardized expected string for initial position expected_start = ( "rnbqkbnr\n" "pppppppp\n" "........\n" "........\n" "........\n" "........\n" "PPPPPPPP\n" "RNBQKBNR" ) # test string representation of board actual_start = str(default_brd) if actual_start != expected_start: raise AssertionError(f"initial board mismatch:\nexpected:\n{expected_start}\n\nactual:\n{actual_start}") # test num moves mvs = default_brd.moves_basic(Color.BLACK) expected_move_count = 20 # 16 pawn moves and 4 knight if len(mvs) != expected_move_count: raise AssertionError(f"initial move count for black mismatch: expected {expected_move_count}, got {len(mvs)}") mvs = default_brd.moves_basic(Color.WHITE) expected_move_count = 20 # 16 pawn moves and 4 knight if len(mvs) != expected_move_count: raise AssertionError(f"Initial move count for white mismatch: expected {expected_move_count}, got {len(mvs)}") play_random_game(verbose=True) print(f"all tests passed") if __name__ == '__main__': main()