simoncreates
315 lines
9.9 KiB
Python
315 lines
9.9 KiB
Python
from enum import Enum, auto
|
|
from dataclasses import dataclass
|
|
from typing import Optional, List, Tuple
|
|
|
|
|
|
class PieceType(Enum):
|
|
PAWN = auto()
|
|
KNIGHT = auto()
|
|
BISHOP = auto()
|
|
ROOK = auto()
|
|
QUEEN = auto()
|
|
KING = auto()
|
|
|
|
|
|
class Color(Enum):
|
|
WHITE = auto()
|
|
BLACK = auto()
|
|
|
|
|
|
@dataclass
|
|
class BoardPos:
|
|
p: Tuple[int, int]
|
|
|
|
def __str__(self) -> str:
|
|
return f"({self.p[0]},{self.p[1]})"
|
|
|
|
|
|
@dataclass
|
|
class Piece:
|
|
type: PieceType
|
|
color: Color
|
|
|
|
def char(self) -> str:
|
|
"""Single-character representation. Uppercase = White, lowercase = Black."""
|
|
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]]
|
|
|
|
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)]
|
|
)
|
|
|
|
# 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
|
|
|
|
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_checked(self, color: Color) -> List[BoardMove]:
|
|
moves: List[BoardMove] = []
|
|
for i, row in enumerate(self.fields):
|
|
for j, field in enumerate(row):
|
|
if field.piece is None:
|
|
continue
|
|
if field.piece.color != color:
|
|
continue
|
|
src_pos = BoardPos((i, j))
|
|
piece = field.piece
|
|
rays = self.moves_unchecked(piece, src_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] != j
|
|
and abs(ray[0][0] - i) == 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(src_pos, BoardPos((tr, tc))))
|
|
continue
|
|
else:
|
|
# occupied
|
|
if target_field.piece.color == piece.color:
|
|
break
|
|
else:
|
|
# opponent piece
|
|
moves.append(BoardMove(src_pos, BoardPos((tr, tc))))
|
|
break
|
|
return moves
|
|
|
|
|
|
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)
|
|
|
|
@classmethod
|
|
def from_string(cls, s: str) -> "ChessBoard":
|
|
lines = s.strip().split("\n")
|
|
|
|
brd = cls([[BoardField() for _ in range(8)] for _ in range(8)])
|
|
|
|
char_to_piece = {
|
|
'p': PieceType.PAWN,
|
|
'n': PieceType.KNIGHT,
|
|
'b': PieceType.BISHOP,
|
|
'r': PieceType.ROOK,
|
|
'q': PieceType.QUEEN,
|
|
'k': PieceType.KING,
|
|
}
|
|
|
|
for r, line in enumerate(lines):
|
|
for c, ch in enumerate(line):
|
|
if ch == '.':
|
|
continue
|
|
lower = ch.lower()
|
|
if lower not in char_to_piece:
|
|
raise ValueError(f"Invalid piece character '{ch}' at ({r},{c})")
|
|
piece_type = char_to_piece[lower]
|
|
color = Color.WHITE if ch.isupper() else Color.BLACK
|
|
brd.place(r, c, piece_type, color)
|
|
|
|
return brd
|
|
|
|
|
|
# used only for testing purposes
|
|
def main():
|
|
default_brd = ChessBoard.init_default()
|
|
|
|
# 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_checked(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_checked(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)}")
|
|
|
|
|
|
# test loading
|
|
loaded_brd = ChessBoard.from_string(expected_start)
|
|
if str(loaded_brd) != expected_start:
|
|
raise AssertionError(
|
|
"board loaded from string does not match original representation"
|
|
)
|
|
|
|
print(f"all tests passed")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|