This post is in succession to this question. I have implemented all functionalities like castling, en-passant, pawn promotion etc. 50 move rule and 3-move repetition is pending.
I would like my code to be reviewed before going to the AI (simple) part of the project. I would appreciate any improvements and also bugs found.
Run the console_ui.py
file to play chess!
White is at the bottom and black is at the top. Uppercase characters represent white pieces and lowercase represent black.
utility.py
from enum import Enum
# Piece and board related stuff
START_MATRIX = [
["R", "N", "B", "Q", "K", "B", "N", "R"],
["P"] * 8,
["."] * 8,
["."] * 8,
["."] * 8,
["."] * 8,
["p"] * 8,
["r", "n", "b", "q", "k", "b", "n", "r"]
]
class Color(Enum):
NONE = 0
WHITE = 1
BLACK = 2
def get_color(piece: str) -> bool:
if piece == ".":
return Color.NONE
elif piece.isupper():
return Color.WHITE
return Color.BLACK
def get_opposite_color(color: Color) -> Color:
if color == Color.WHITE:
return Color.BLACK
elif color == Color.BLACK:
return Color.WHITE
return color.NONE
def is_white(piece: str) -> bool:
return get_color(piece) == Color.WHITE
def is_black(piece: str) -> bool:
return get_color(piece) == Color.BLACK
def is_empty(piece: str) -> bool:
return piece == "."
def is_king(piece: str) -> bool:
return piece.lower() == "k"
def is_queen(piece: str) -> bool:
return piece.lower() == "q"
def is_rook(piece: str) -> bool:
return piece.lower() == "r"
def is_bishop(piece: str) -> bool:
return piece.lower() == "b"
def is_knight(piece: str) -> bool:
return piece.lower() == "n"
def is_pawn(piece: str) -> bool:
return piece.lower() == "p"
# Piece movement related stuff
def is_in_bounds(*nums: int):
return all([0 <= num <= 7 for num in nums])
class MoveType(Enum):
NORMAL = 0
SHORT_CASTLE = 1
LONG_CASTLE = 2
PAWN_PROMOTION = 3
FILE_TO_INDEX_MAP = {
"a": 0, "b": 1, "c": 2, "d": 3,
"e": 4, "f": 5, "g": 6, "h": 7
}
INDEX_TO_FILE_MAP = {v: k for k, v in FILE_TO_INDEX_MAP.items()}
def encode_move(move: dict) -> str:
# Normal move: e2e4
if move["type"] == MoveType.NORMAL:
start_y, start_x = move["start_pos"]
end_y, end_x = move["end_pos"]
return (
INDEX_TO_FILE_MAP[start_x] + str(start_y + 1)
+ INDEX_TO_FILE_MAP[end_x] + str(end_y + 1)
)
elif move["type"] == MoveType.PAWN_PROMOTION:
start_y, start_x = move["start_pos"]
end_y, end_x = move["end_pos"]
return (
INDEX_TO_FILE_MAP[start_x] + str(start_y + 1)
+ INDEX_TO_FILE_MAP[end_x] + str(end_y + 1)
+ "=" + move["promoted_piece"]
)
# Short castle: 0-0
elif move["type"] == MoveType.SHORT_CASTLE:
return "0-0"
# Long castle: 0-0-0
elif move["type"] == MoveType.LONG_CASTLE:
return "0-0-0"
# Pawn promotion: ToDo
def decode_move(move: str) -> dict:
# Normal move: e2e4
if len(move) == 4:
return {
"type": MoveType.NORMAL,
"start_pos": [int(move[1]) - 1, FILE_TO_INDEX_MAP[move[0]]],
"end_pos": [int(move[3]) - 1, FILE_TO_INDEX_MAP[move[2]]],
}
# Pawn Promotion: e7e8=q
elif len(move) == 6:
return {
"type": MoveType.PAWN_PROMOTION,
"start_pos": [int(move[1]) - 1, FILE_TO_INDEX_MAP[move[0]]],
"end_pos": [int(move[3]) - 1, FILE_TO_INDEX_MAP[move[2]]],
"promoted_piece": move[5]
}
# Short castle: 0-0
elif move == "0-0":
return {
"type": MoveType.SHORT_CASTLE
}
# Long castle: 0-0-0
elif move == "0-0-0":
return {
"type": MoveType.LONG_CASTLE
}
# Pawn promotion: ToDo
KNIGHT_DELTAS = [
# delta_y, delta_x
[2, 1], [1, 2],
[-2, 1], [-1, 2],
[1, -2], [2, -1],
[-1, -2], [-2, -1]
]
BISHOP_DELTAS = [
# delta_y, delta_x
[1, 1], [-1, -1],
[-1, 1], [1, -1]
]
ROOK_DELTAS = [
# delta_y, delta_x
[1, 0], [0, 1],
[-1, 0], [0, -1]
]
QUEEN_DELTAS = BISHOP_DELTAS + ROOK_DELTAS
KING_DELTAS = [
# delta_y, delta_x
[1, 0], [0, 1],
[-1, 0], [0, -1],
[1, 1], [-1, -1],
[-1, 1], [1, -1]
]
board.py
from utility import *
from copy import deepcopy
class GameState(Enum):
RUNNING = 0
CHECKMATE = 1
STALEMATE = 2
class Board:
def __init__(self, matrix=None, current_turn_color=Color.WHITE) -> None:
self.matrix = matrix
if not matrix:
self.matrix = START_MATRIX
self.current_turn_color = current_turn_color
self.game_state = GameState.RUNNING
self.moves_history = []
self.en_passant_squares = {
Color.WHITE: [],
Color.BLACK: []
}
self.can_castle = {
Color.WHITE: {
MoveType.SHORT_CASTLE: True,
MoveType.LONG_CASTLE: True
},
Color.BLACK: {
MoveType.SHORT_CASTLE: True,
MoveType.LONG_CASTLE: True
}
}
self.legal_moves = self.generate_legal_moves()
def play_move(self, str_move: str) -> None:
if str_move in self.legal_moves:
move_data = decode_move(str_move)
print("Printing en-passant squares: ", self.en_passant_squares)
self.implement_move(move_data)
self.en_passant_squares[self.current_turn_color] = []
self.current_turn_color = get_opposite_color(self.current_turn_color)
self.legal_moves = self.generate_legal_moves()
self.moves_history.append(str_move)
self.update_game_state()
else:
print(f"{str_move} not legal!! Error!!")
return
def update_game_state(self):
if len(self.legal_moves) == 0:
if self.is_check():
self.game_state = GameState.CHECKMATE
else:
self.game_state = GameState.STALEMATE
else:
self.game_state = GameState.RUNNING
def implement_move(self, move: dict) -> None:
if move["type"] == MoveType.NORMAL:
start_y, start_x = move["start_pos"]
end_y, end_x = move["end_pos"]
piece = self.matrix[start_y][start_x]
self.matrix[end_y][end_x] = piece
self.matrix[start_y][start_x] = "."
# Prevent castle
if is_rook(piece):
if start_x == 7: # h1/h8 rook
self.can_castle[self.current_turn_color][MoveType.SHORT_CASTLE] = False
elif start_x == 0: # a1/a8 rook
self.can_castle[self.current_turn_color][MoveType.LONG_CASTLE] = False
# En passant
elif is_pawn(piece):
# If en-passant happens
if [end_y, end_x] in self.en_passant_squares[self.current_turn_color]:
self.matrix[start_y][end_x] = "."
# If black pawn moves 2 steps ahead
if start_y - end_y == 2:
self.en_passant_squares[Color.WHITE].append([start_y - 1, start_x])
# If white pawn moves 2 steps ahead
elif start_y - end_y == -2:
self.en_passant_squares[Color.BLACK].append([start_y + 1, start_x])
elif move["type"] == MoveType.PAWN_PROMOTION:
start_y, start_x = move["start_pos"]
end_y, end_x = move["end_pos"]
promoted_piece = move["promoted_piece"]
if self.current_turn_color == Color.WHITE:
promoted_piece = promoted_piece.upper()
self.matrix[end_y][end_x] = promoted_piece
self.matrix[start_y][start_x] = "."
elif move["type"] == MoveType.SHORT_CASTLE:
pos_y = 0 if self.current_turn_color == Color.WHITE else 7
# Move h1/h8 rook to f1/f8
self.matrix[pos_y][5] = self.matrix[pos_y][7]
self.matrix[pos_y][7] = "."
# Move e1/e8 king to g1/g8
self.matrix[pos_y][6] = self.matrix[pos_y][4]
self.matrix[pos_y][4] = "."
self.can_castle[self.current_turn_color][MoveType.SHORT_CASTLE] = False
self.can_castle[self.current_turn_color][MoveType.LONG_CASTLE] = False
elif move["type"] == MoveType.LONG_CASTLE:
pos_y = 0 if self.current_turn_color == Color.WHITE else 7
# Move a1/a8 rook to d1/d8
self.matrix[pos_y][3] = self.matrix[pos_y][0]
self.matrix[pos_y][0] = "."
# Move e1/e8 king to c1/c8
self.matrix[pos_y][2] = self.matrix[pos_y][4]
self.matrix[pos_y][4] = "."
self.can_castle[self.current_turn_color][MoveType.SHORT_CASTLE] = False
self.can_castle[self.current_turn_color][MoveType.LONG_CASTLE] = False
def get_all_pieces_pos(self, color: Color) -> list:
pieces_pos = []
for y, rank in enumerate(self.matrix):
for x, piece in enumerate(rank):
if type(piece) == int:
print(rank, x, piece)
if get_color(piece) == color:
pieces_pos.append([y, x])
return pieces_pos
def generate_primitive_moves(self) -> list[str]:
"Generates all move without validating any checks"
primitive_moves = []
# Normal moves
pieces_pos = self.get_all_pieces_pos(self.current_turn_color)
opposite_color = get_opposite_color(self.current_turn_color)
# print("pieces_pos:", pieces_pos)
for pos_y, pos_x in pieces_pos:
piece = self.matrix[pos_y][pos_x]
# print(f"pos_y: {pos_y}, pos_x: {pos_x}, piece: {piece}")
if is_pawn(piece):
delta_y = 1 if self.current_turn_color == Color.WHITE else -1
# Forward 1 square
if is_in_bounds(pos_y + delta_y) and is_empty(self.matrix[pos_y + delta_y][pos_x]):
if (pos_y + delta_y) in [0, 7]:
for promoted_piece in ["q", "r", "n", "b"]:
primitive_moves.append(encode_move({
"type": MoveType.PAWN_PROMOTION,
"start_pos": [pos_y, pos_x],
"end_pos": [pos_y + delta_y, pos_x],
"promoted_piece": promoted_piece
}))
else:
primitive_moves.append(encode_move({
"type": MoveType.NORMAL,
"start_pos": [pos_y, pos_x],
"end_pos": [pos_y + delta_y, pos_x]
}))
# Forward 2 square
if is_in_bounds(pos_y + 2*delta_y) and is_empty(self.matrix[pos_y + 2*delta_y][pos_x]):
if (
(self.current_turn_color == Color.WHITE and pos_y == 1) or
(self.current_turn_color == Color.BLACK and pos_y == 6)
):
primitive_moves.append(encode_move({
"type": MoveType.NORMAL,
"start_pos": [pos_y, pos_x],
"end_pos": [pos_y + 2*delta_y, pos_x]
}))
# Diagonal 1 square left/right
for delta_x in [1, -1]:
if (
is_in_bounds(pos_y + delta_y, pos_x + delta_x) and
(
get_color(self.matrix[pos_y + delta_y][pos_x + delta_x]) == opposite_color or
[pos_y + delta_y, pos_x + delta_x] in self.en_passant_squares[self.current_turn_color]
)
):
if (pos_y + delta_y) in [0, 7]:
for promoted_piece in ["q", "r", "n", "b"]:
primitive_moves.append(encode_move({
"type": MoveType.PAWN_PROMOTION,
"start_pos": [pos_y, pos_x],
"end_pos": [pos_y + delta_y, pos_x + delta_x],
"promoted_piece": promoted_piece
}))
else:
primitive_moves.append(encode_move({
"type": MoveType.NORMAL,
"start_pos": [pos_y, pos_x],
"end_pos": [pos_y + delta_y, pos_x + delta_x]
}))
elif is_knight(piece):
primitive_moves += self.generate_knight_legal_moves(pos_y, pos_x)
elif is_bishop(piece):
primitive_moves += self.generate_bishop_legal_moves(pos_y, pos_x)
elif is_rook(piece):
primitive_moves += self.generate_rook_legal_moves(pos_y, pos_x)
elif is_queen(piece):
primitive_moves += self.generate_queen_legal_moves(pos_y, pos_x)
elif is_king(piece):
primitive_moves += self.generate_king_legal_moves(pos_y, pos_x)
pos_y = 0 if self.current_turn_color == Color.WHITE else 7
if (
self.can_castle[self.current_turn_color][MoveType.SHORT_CASTLE] and
is_empty(self.matrix[pos_y][5]) and
is_empty(self.matrix[pos_y][6])
):
primitive_moves.append("0-0")
if (
self.can_castle[self.current_turn_color][MoveType.LONG_CASTLE] and
is_empty(self.matrix[pos_y][1]) and
is_empty(self.matrix[pos_y][2]) and
is_empty(self.matrix[pos_y][3])
):
primitive_moves.append("0-0-0")
return primitive_moves
def generate_legal_moves(self) -> list[str]:
legal_moves = []
for move in self.generate_primitive_moves():
copy_board = deepcopy(self)
copy_board.implement_move(decode_move(move))
if not copy_board.is_check():
if move == "0-0":
start_y = 0 if self.current_turn_color == Color.WHITE else 7
if not copy_board.is_attacked(start_y, 5, get_opposite_color(self.current_turn_color)):
legal_moves.append(move)
elif move == "0-0-0":
start_y = 0 if self.current_turn_color == Color.WHITE else 7
if not copy_board.is_attacked(start_y, 3, get_opposite_color(self.current_turn_color)):
legal_moves.append(move)
else:
legal_moves.append(move)
return legal_moves
def is_attacked(self, pos_y: int, pos_x: int, color: Color) -> list:
original_turn_color = self.current_turn_color
self.current_turn_color = color
for move in self.generate_primitive_moves():
move_data = decode_move(move)
if move_data["type"] == MoveType.NORMAL:
start_y, start_x = move_data["start_pos"]
if move_data["end_pos"] == [pos_y, pos_x]:
if is_pawn(self.matrix[start_y][start_x]):
if abs(start_x - pos_x) == 1:
self.current_turn_color = original_turn_color
return True
else:
self.current_turn_color = original_turn_color
return True
self.current_turn_color = original_turn_color
return False
def find_king(self, color: Color) -> list[int, int]:
for y, rank in enumerate(self.matrix):
for x, piece in enumerate(rank):
if is_king(piece) and get_color(piece) == color:
return [y, x]
return False
def is_check(self):
king_pos = self.find_king(self.current_turn_color)
return self.is_attacked(*king_pos, get_opposite_color(self.current_turn_color))
def generate_knight_legal_moves(self, pos_y: int, pos_x: int) -> list[str]:
legal_moves = []
for delta_y, delta_x in KNIGHT_DELTAS:
end_y = pos_y + delta_y
end_x = pos_x + delta_x
if (
is_in_bounds(end_y, end_x) and
get_color(self.matrix[end_y][end_x]) != self.current_turn_color
):
legal_moves.append(encode_move({
"type": MoveType.NORMAL,
"start_pos": [pos_y, pos_x],
"end_pos": [end_y, end_x]
}))
return legal_moves
def generate_bishop_legal_moves(self, pos_y: int, pos_x: int) -> list[str]:
legal_moves = []
for delta_y, delta_x in BISHOP_DELTAS:
for i in range(1, 9):
end_y = pos_y + i*delta_y
end_x = pos_x + i*delta_x
if (
not is_in_bounds(end_y, end_x) or
get_color(self.matrix[end_y][end_x]) == self.current_turn_color
):
break
legal_moves.append(encode_move({
"type": MoveType.NORMAL,
"start_pos": [pos_y, pos_x],
"end_pos": [end_y, end_x]
}))
if not is_empty(self.matrix[end_y][end_x]):
break
return legal_moves
def generate_rook_legal_moves(self, pos_y: int, pos_x: int) -> list[str]:
legal_moves = []
for delta_y, delta_x in ROOK_DELTAS:
for i in range(1, 9):
end_y = pos_y + i*delta_y
end_x = pos_x + i*delta_x
if (
not is_in_bounds(end_y, end_x) or
get_color(self.matrix[end_y][end_x]) == self.current_turn_color
):
break
legal_moves.append(encode_move({
"type": MoveType.NORMAL,
"start_pos": [pos_y, pos_x],
"end_pos": [end_y, end_x]
}))
if not is_empty(self.matrix[end_y][end_x]):
break
return legal_moves
def generate_queen_legal_moves(self, pos_y: int, pos_x: int) -> list[str]:
return (
self.generate_bishop_legal_moves(pos_y, pos_x) +
self.generate_rook_legal_moves(pos_y, pos_x)
)
def generate_king_legal_moves(self, pos_y: int, pos_x: int) -> list[str]:
legal_moves = []
for delta_y, delta_x in KING_DELTAS:
end_y = pos_y + delta_y
end_x = pos_x + delta_x
if (
is_in_bounds(end_y, end_x) and
get_color(self.matrix[end_y][end_x]) != self.current_turn_color
):
legal_moves.append(encode_move({
"type": MoveType.NORMAL,
"start_pos": [pos_y, pos_x],
"end_pos": [end_y, end_x]
}))
return legal_moves
def __repr__(self) -> str:
output = ""
for row in reversed(self.matrix):
output += " ".join(row) + "\n"
return output[:-1]
console_ui.py
from board import *
game = Board()
def print_chess_board():
new_array = []
reversed_matrix = list(reversed(game.matrix))
for i in range(8):
new_array.append([str(i + 1), *reversed_matrix[i]])
new_array.append([" ", "a", "b", "c", "d", "e", "f", "g", "h"])
print("\n".join([" ".join(rank) for rank in new_array]))
while True:
print_chess_board()
if game.game_state == GameState.CHECKMATE:
print("Checkmate!!")
print(f"{turn} wins the game!!")
break
elif game.game_state == GameState.STALEMATE:
print("Stalemate!! The game is a draw!")
break
turn = "White" if game.current_turn_color == Color.WHITE else "Black"
move = input(f"{turn}!! Enter your move: ").lower().strip()
game.play_move(move)
Edit 1: There is small mistake in the interface, the numbering from 1 to 8 is written in inverse order.
b1c3
as first move as white, it actually moves black's knight from b8 to c6. \$\endgroup\$