checkmate_checker.py

"""
Determine if a given chess position is "checkmate" or not.

Input is a FEN string, a commonly used way to 'serialize' the state of a chess game at a certain move,
it looks like this:

rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1

Explanation of algorithm:

Evaluate if there is a check on the board, if no check return false

If there is check, make all possible moves for the side in check and after each
one evaluate if there is a check on the board.

If after each possible move there is still a check, it is checkmate, return true.
"""
import copy
import string

from enum import Enum
from typing import Optional, List, Tuple

BOARD_DIMENSION = 8


class PieceType(Enum):
    PAWN = 2
    KNIGHT = 3
    BISHOP = 4
    ROOK = 5
    QUEEN = 6
    KING = 7


class Color(Enum):
    WHITE = 1
    BLACK = 2


SYMBOL_TO_PIECE = {
    'p': PieceType.PAWN,
    'n': PieceType.KNIGHT,
    'b': PieceType.BISHOP,
    'r': PieceType.ROOK,
    'q': PieceType.QUEEN,
    'k': PieceType.KING
}

PIECE_TO_SYMBOL = {v: k for (k, v) in SYMBOL_TO_PIECE.items()}

PAWN_MOVEMENTS = {
    Color.BLACK: [[1, 0]],
    Color.WHITE: [[-1, 0]]
}

PAWN_CAPTURE_MOVEMENTS = {
    Color.BLACK: [[1, 1], [1, -1]],
    Color.WHITE: [[-1, 1], [-1, -1]]
}

BISHOP_MOVEMENTS = [
    [1, 1], [1, -1], [-1, 1], [-1, -1]
]

ROOK_MOVEMENTS = [
    [1, 0], [0, 1], [-1, 0], [0, -1]
]

QUEEN_AND_KING_MOVEMENTS = ROOK_MOVEMENTS + BISHOP_MOVEMENTS

KNIGHT_MOVEMENTS = [
    [2, 1], [1, 2], [-2, 1], [-1, 2],
    [2, -1], [1, -2], [-2, -1], [-1, -2]
]

MOVEMENTS = {
    PieceType.ROOK: ROOK_MOVEMENTS,
    PieceType.BISHOP: BISHOP_MOVEMENTS,
    PieceType.KING: QUEEN_AND_KING_MOVEMENTS,
    PieceType.QUEEN: QUEEN_AND_KING_MOVEMENTS,
    PieceType.PAWN: PAWN_MOVEMENTS,
    PieceType.KNIGHT: KNIGHT_MOVEMENTS
}


def sum_vector(a: Tuple[int, int], b: Tuple[int, int]):
    rank_a, file_a = a
    rank_b, file_b = b
    return rank_a + rank_b, file_a + file_b


def multiply_vector(v: Tuple[int, int], factor: int):
    rank, file = v
    return rank * factor, file * factor


def coordinate_is_on_board(rank: int, file: int, board_dimension: int):
    return 0 <= rank < board_dimension and 0 <= file < board_dimension


class SquareState:
    def __init__(self, piece_type: PieceType, piece_color: Color):
        self.piece_type = piece_type
        self.piece_color = piece_color
        self.attacked = None

    def is_empty(self):
        return self.piece_type is None

    def get_color(self):
        return self.piece_color

    def get_piece_type(self):
        return self.piece_type

    def __str__(self):
        if self.piece_type is None:
            return '.'

        symbol = PIECE_TO_SYMBOL[self.piece_type]
        if self.piece_color == Color.WHITE:
            return symbol.upper()

        return symbol.lower()

    def __repr__(self):
        return self.__str__()


class Board:
    def __init__(self, board_dimension: int, to_move_color: Color, en_passant_square: Optional[Tuple[int, int]],
                 data=None):
        self.board_dimension = board_dimension
        self.en_passant_square = en_passant_square
        self.data = [[None] * board_dimension for _ in range(board_dimension)] if data is None else data
        self.to_move_color = to_move_color

    def set_square(self, rank, file, state: SquareState):
        self.data[rank][file] = state

    def get_square(self, rank, file):
        return self.data[rank][file]

    def get_dimension(self) -> int:
        return self.board_dimension

    def has_en_passant_square(self) -> bool:
        return self.en_passant_square is not None

    def get_en_passant_square(self) -> (int, int):
        return self.en_passant_square

    def get_to_move_color(self):
        return self.to_move_color

    def is_en_passant_capture(self, coords_a, coords_b):
        """
        If a pawn changes files, and the square its moving to is empty,
        then it must be performing an en passant capture.
        """
        rank_a, file_a = coords_a
        rank_b, file_b = coords_b
        return (self.data[rank_b][file_b].is_empty() and
                self.data[rank_a][file_a].get_piece_type() == PieceType.PAWN and
                file_a != file_b)

    def move(self, coords_a, coords_b):
        rank_a, file_a = coords_a
        rank_b, file_b = coords_b

        # Normally the piece we capture is on the square we move to. With
        # en passant however, the piece we capture is NOT on the square we
        # moved to, so we need extra logic to make sure the captured piece
        # is taken off the board. We know the captured pawn is right behind
        # where we moved to, so we look back one rank and remove that pawn.
        if self.is_en_passant_capture(coords_a, coords_b):
            if self.to_move_color == Color.WHITE:
                self.data[rank_b + 1][file_b] = SquareState(None, None)
            else:
                self.data[rank_b - 1][file_b] = SquareState(None, None)

        self.data[rank_b][file_b] = self.data[rank_a][file_a]
        self.data[rank_a][file_a] = SquareState(None, None)

    def pieces(self):
        for rank in range(BOARD_DIMENSION):
            for file in range(BOARD_DIMENSION):
                square_state = self.data[rank][file]
                if not square_state.is_empty():
                    yield square_state, (rank, file)

    def king_position(self, color):
        for rank in range(BOARD_DIMENSION):
            for file in range(BOARD_DIMENSION):
                square_state = self.data[rank][file]
                if (not square_state.is_empty() and
                        square_state.get_piece_type() == PieceType.KING and
                        square_state.get_color() == color):
                    return rank, file

    def copy(self):
        return Board(
            board_dimension=self.get_dimension(),
            to_move_color=self.get_to_move_color(),
            en_passant_square=self.get_en_passant_square(),
            data=copy.deepcopy(self.data))

    def __str__(self):
        return '\n'.join([''.join([str(x) for x in rank]) for rank in self.data])

    def __repr__(self):
        return self.__str__()


def pawn_is_on_starting_square(color: Color, rank: int, board_dimension):
    """
    Pawns are incapable of ever moving laterally, so the rank of the pawn and its color
    are sufficient to tell if it is on its starting square.
    """
    if color == Color.BLACK and rank == 1:
        return True

    if color == Color.WHITE and rank == board_dimension - 2:
        return True

    return False


def candidate_moves_pawn(rank: int, file: int, color: Color, board: Board, checking_moves_only: bool):
    """
    Pawns have four different modes of movement, and so a special function is needed to cover this complexity.

    Mode 1: A single move forward.
    Mode 2: A double move forward, only possible on that pawn's first move
    Mode 3: A capture of another piece diagonally forward one square
    Mode 4: "En passant" capture, where a pawn can capture a pawn laterally
    if that other pawn moved with Mode 2 last turn.

    A pawn can only give check using Mode 3, so we allow the other modes to be toggled off with the
    "attacking_moves_only" flag.
    """
    is_on_starting_sqaure = pawn_is_on_starting_square(
        color=color,
        rank=rank,
        board_dimension=board.get_dimension())

    for movement in PAWN_CAPTURE_MOVEMENTS[color]:
        seen_rank, seen_file = sum_vector(movement, (rank, file))
        if (coordinate_is_on_board(seen_rank, seen_file, board.get_dimension()) and
                not board.get_square(seen_rank, seen_file).is_empty() and
                board.get_square(seen_rank, seen_file).get_color() != color):
            yield seen_rank, seen_file

    if not checking_moves_only:
        if board.has_en_passant_square() and (seen_rank, seen_file) == board.get_en_passant_square():
            yield seen_rank, seen_file

        for new_rank, new_file in candidate_moves_of_linear_mover(
                rank=rank,
                file=file,
                color=color,
                movements=PAWN_MOVEMENTS[color],
                board=board,
                max_moves=(2 if is_on_starting_sqaure else 1),
                can_capture_in_line=False):
            yield new_rank, new_file


def candidate_moves_of_linear_mover(
        rank: int,
        file: int,
        color: Color,
        movements: List[Tuple[int, int]],
        board: Board,
        max_moves: int,
        can_capture_in_line: bool):
    """
    This function covers the movement of rooks, bishops, queens, kings, and pawn movement modes 1 and 2.

    All of these pieces move linearly 1 or more spaces, the max_moves arguments specifies
    how many spaces the piece can move at a time. All these pieces also cannot move "through"
    another piece, which is unique to the knight.
    """
    for movement in movements:
        for factor in range(1, max_moves + 1):
            new_vector = multiply_vector(movement, factor)
            seen_rank, seen_file = sum_vector(new_vector, (rank, file))

            if not coordinate_is_on_board(seen_rank, seen_file, board.get_dimension()):
                break

            if not board.get_square(seen_rank, seen_file).is_empty():
                if can_capture_in_line and board.get_square(seen_rank, seen_file).get_color() != color:
                    yield seen_rank, seen_file
                break
            else:
                yield seen_rank, seen_file


def candidate_moves_knight(rank, file, color, board):
    """
    Knights can "jump" over pieces and so a special case is needed to cover their movements.
    """
    for movement in KNIGHT_MOVEMENTS:
        seen_rank, seen_file = sum_vector(movement, (rank, file))

        if not coordinate_is_on_board(seen_rank, seen_file, board.get_dimension()):
            continue

        if (not board.get_square(seen_rank, seen_file).is_empty()
                and board.get_square(seen_rank, seen_file).get_color() == color):
            continue

        yield seen_rank, seen_file


def get_candidate_move_squares(piece_type: PieceType,
                               color: Color,
                               rank: int,
                               file: int,
                               board: Board,
                               checking_moves_only: bool):
    """
    Get all possible squares that a piece could move to this turn.

    Note: we have no logic here for "castling" a special move where the
    king and rook exchange positions. Castling is irrelevant for determining
    checkmate because it is illegal to castle when the king is in check.
    """
    if piece_type == PieceType.PAWN:
        return candidate_moves_pawn(rank=rank,
                                    file=file,
                                    color=color,
                                    board=board,
                                    checking_moves_only=checking_moves_only)
    elif piece_type == PieceType.KNIGHT:
        return candidate_moves_knight(rank=rank,
                                      file=file,
                                      color=color,
                                      board=board)
    elif piece_type in (PieceType.KING, PieceType.ROOK, PieceType.BISHOP, PieceType.QUEEN):
        return candidate_moves_of_linear_mover(
            rank=rank,
            file=file,
            color=color,
            movements=MOVEMENTS[piece_type],
            board=board,
            max_moves=1 if piece_type == PieceType.KING else board.get_dimension(),
            can_capture_in_line=True)


def is_check(board: Board):
    king_rank, king_file = board.king_position(board.get_to_move_color())

    for piece, (rank, file) in board.pieces():
        if piece.get_color() != board.get_to_move_color():
            for new_rank, new_file in get_candidate_move_squares(
                    piece_type=piece.get_piece_type(),
                    color=piece.get_color(),
                    rank=rank,
                    file=file,
                    board=board,
                    checking_moves_only=True):
                if (new_rank, new_file) == (king_rank, king_file):
                    return True
    return False


def is_check_mate(board: Board):
    if not is_check(board):
        return False

    for piece, (rank, file) in board.pieces():
        if piece.get_color() == board.get_to_move_color():
            for new_rank, new_file in get_candidate_move_squares(
                    piece_type=piece.get_piece_type(),
                    color=piece.get_color(),
                    rank=rank,
                    file=file,
                    board=board,
                    checking_moves_only=False):
                # We make a copy of the board so we can return to the
                # previous state before we moved this piece.
                possible_board = board.copy()
                possible_board.move((rank, file), (new_rank, new_file))

                if not is_check(possible_board):
                    return False

    return True


def parse_en_passant_sqaure(coordinate):
    """
    Convert algebraic chess square notation to coordinates:

    a8 => (0, 0)
    h1 => (7, 7)
    """
    return string.ascii_letters.index(coordinate[0]), 7 - (int(coordinate[1]) - 1)


def parse_fen(fen_string: str) -> Board:
    position, to_move, castling_availabilty, en_passant_target_square, half_move_clock, move_count = fen_string.split(
        ' ')

    board = Board(
        board_dimension=BOARD_DIMENSION,
        to_move_color=Color.WHITE if to_move == 'w' else Color.BLACK,
        en_passant_square=parse_en_passant_sqaure(
            en_passant_target_square) if en_passant_target_square != '-' else None)

    for (rank, rank_position) in enumerate(position.split('/')):
        file = 0
        for square_position in rank_position:
            if square_position in string.digits:
                for i in range(int(square_position)):
                    board.set_square(rank, file, SquareState(None, None))
                    file += 1
            elif square_position.lower() in SYMBOL_TO_PIECE:
                color = Color.WHITE if square_position.isupper() else Color.BLACK
                piece_type = SYMBOL_TO_PIECE[square_position.lower()]
                board.set_square(rank, file, SquareState(piece_type, color))
                file += 1

    return board


def main():
    while True:
        fen_string = input('Paste FEN: ').strip()
        print('Checkmate!' if is_check_mate(parse_fen(fen_string)) else 'Not checkmate..')


if __name__ == '__main__':
    main()

checkmate_checker_test.py

from checkmate_checker import is_check_mate, parse_fen

def test_main():
    for fen_string in [
        '4rrk1/p4ppp/1p6/2pP1Q2/2BK2P1/3P4/P1q2P1P/R1R5 w - - 0 1',
        '4k3/4PP2/2P3P1/8/4Q3/8/4K3/8 b - - 0 1',
        '4k2R/8/4K3/8/8/8/8/8 b - - 0 1',
        '4k3/7N/2N1K1B1/8/8/8/8/8 b - - 0 1',
        '4k3/4Q3/4K3/8/8/8/8/8 b - - 0 1',
        '4k3/4PP2/4K3/8/8/8/8/8 b - - 0 1',
        '8/8/8/8/8/4k3/8/r3K3 w - - 0 1',
        '8/8/8/1k6/4K3/8/2P5/3rrr2 w - - 0 1'
    ]:
        assert is_check_mate(parse_fen(fen_string))

    for fen_string in [
        '4rrk1/p4ppp/1p6/2pP1Q2/2BK2P1/3P4/P1q2P1P/R1R5 w - c6 0 1',
        '4k3/4P3/2P3P1/8/4Q3/8/4K3/8 b - - 0 1',
        '4k2R/8/4K3/5r2/8/8/8/8 b - - 0 1',
        '3k3R/8/4K3/8/8/8/8/8 b - - 0 1',
        '4k2R/8/4K2b/8/8/8/8/8 b - - 0 1',
        '3k3R/8/3K2n1/8/8/8/8/8 b - - 0 1',
        '3k3R/8/3K4/7r/8/8/8/8 b - - 0 1',
        '3k3R/8/3K4/8/8/8/8/q7 b - - 0 1',
        'R7/2p5/7R/1k5R/8/1K6/8/8 b - - 0 1',
        '1BQ5/2p5/1k5R/7R/1K6/8/8/8 b - - 0 1',
        '8/8/1k6/8/1K5r/7r/2P5/8 w - - 0 1',
        '8/8/8/1k6/8/1K5r/2P4r/8 w - - 0 1'
    ]:
        assert not is_check_mate(parse_fen(fen_string))