peanut chess engine

constants.py

BOARD_SIZE = 8
FIRST_RANK = 1
LAST_RANK = BOARD_SIZE

PAWN_START = 2

KING_START_FILE = 5

KNIGHT, KING, BISHOP, ROOK, QUEEN, PAWN = (
    'N', 'K', 'B', 'R', 'Q', 'P'
)

CASTLEING = 'C'

WHITE, BLACK = True, False
KINGSIDE, QUEENSIDE = True, False

COLOUR_INDEX, PIECE_INDEX = range(2)


MOVE_PIECE_NAME, MOVE_X, MOVE_Y, MOVE_DEST_X, MOVE_DEST_Y, MOVE_FLAGS = \
range(6)

initial_position.py

from constants import (
    WHITE, BLACK,
    PAWN, ROOK, KNIGHT, BISHOP, KING, QUEEN,
    KING_START_FILE,
    KINGSIDE, QUEENSIDE, 
    PAWN_START, 
    BOARD_SIZE, 
    FIRST_RANK, LAST_RANK
)

INITIAL_POSITION = {}

# setup the pawns

INITIAL_POSITION.update(
    ( (x, PAWN_START), (WHITE, PAWN) )
    for x in range(1, BOARD_SIZE+1)
)
INITIAL_POSITION.update(
    ( (x, BOARD_SIZE-PAWN_START), (BLACK, PAWN) )
    for x in range(1, BOARD_SIZE+1)
)

# setup the rooks, knights, and bishops
for piece, x in ( (ROOK, 1), (KNIGHT, 2), (BISHOP, 3) ):
    INITIAL_POSITION.update((
        ( (x, FIRST_RANK), (WHITE, piece) ),
        ( (BOARD_SIZE-x+1, FIRST_RANK), (WHITE, piece) ),
        ( (x, LAST_RANK), (BLACK, piece) ),
        ( (BOARD_SIZE-x+1, LAST_RANK), (BLACK, piece) )
    ))

# setup the kings and queens
INITIAL_POSITION.update( 
    ( (KING_START_FILE+d, y), (white_turn, piece) )
    for (piece, d) in ( (KING, 0), (QUEEN, -1) )
    for (white_turn, y) in ( (WHITE, FIRST_RANK), (BLACK, LAST_RANK) )
)
    
# castling possible
INITIAL_POSITION[(WHITE, KINGSIDE)] = (True, None)
INITIAL_POSITION[(WHITE, QUEENSIDE)] = (True, None)
INITIAL_POSITION[(BLACK, KINGSIDE)] = (True, None)
INITIAL_POSITION[(BLACK, QUEENSIDE)] = (True, None)

move_generation.py

from constants import (
    FIRST_RANK, LAST_RANK, BOARD_SIZE, 
    PAWN_START, KING_START_FILE,
    KNIGHT, KING, BISHOP, ROOK, QUEEN, PAWN,
    CASTLEING, 
    KINGSIDE, QUEENSIDE,
    COLOUR_INDEX, PIECE_INDEX,
    MOVE_PIECE_NAME, MOVE_X, MOVE_Y, MOVE_DEST_X, MOVE_DEST_Y, MOVE_FLAGS
    )                                                                   

from itertools import chain

def range_excluding_zero(start, end):
    return (i
            for i in range(start, end) 
            if i != 0 )

PIECE_MOVES = {
    KNIGHT: tuple(
        (x, y * ( 1 if abs(x)==2 else 2  ) )
        for x in (-2, -1, 1, 2)
        for y in (1, -1)
    ),

    KING: tuple(
        (x, y)
        for x in (-1,0,1)
        for y in (1,0,1)
        if not (x==0 and y==0)
    ),

    BISHOP: tuple( (x, x*y)
                   # -7 to 7 inclusive
                   for x in range_excluding_zero(-BOARD_SIZE+1,BOARD_SIZE)
                   for y in (1, -1)
               ),

    ROOK: ( 
        tuple( (x,0)
               for x in range_excluding_zero(-BOARD_SIZE+1,BOARD_SIZE) ) +
        tuple( (0,y)
               for y in range_excluding_zero(-BOARD_SIZE+1,BOARD_SIZE) )
    )
}

PIECE_MOVES[QUEEN] = PIECE_MOVES[BISHOP] + PIECE_MOVES[ROOK]

def is_on_board(x, y):
    return 0 < x <= BOARD_SIZE and 0< y <=BOARD_SIZE

def neg_one_zero_or_pos_one(n):
    return 0 if n == 0 else n//abs(n)

def empty_squares_in_between(position, x,y, dest_x, dest_y):
    dir_x = neg_one_zero_or_pos_one(dest_x-x)
    dir_y = neg_one_zero_or_pos_one(dest_y-y)
    while True:
        x += dir_x
        y += dir_y
        if x == dest_x and y == dest_y:
            return True
        elif (x, y) in position:
            return False
            
def piece_of_same_color_on_square(position, white_turn, test_x, test_y):
    cords = (test_x, test_y)
    return cords in position and position[cords][COLOUR_INDEX] == white_turn

def generate_possible_piece_moves(position, white_turn):
    return (
        (piece_name, x, y, x+rel_x, y+rel_y, None)
        for (x, y), (piece_is_white, piece_name) in position.items()
        if piece_name not in (PAWN, None)
        for rel_x, rel_y in PIECE_MOVES[piece_name]
        if piece_is_white == white_turn
        if is_on_board(x+rel_x, y+rel_y)
        if (piece_name == KNIGHT or piece_name == KING or 
            empty_squares_in_between(position, x, y, x+rel_x, y+rel_y) )
        if not piece_of_same_color_on_square(position, piece_is_white, 
                                             x+rel_x, y+rel_y)
    )

def pawns_of_turn(position, white_turn):
    return ( (x, y)
             for (x, y), (piece_is_white, piece_name) in position.items()
             if piece_is_white == white_turn and piece_name == PAWN
         )

def move_pawn_forward(x, y, piece_is_white, n=1):
    return (PAWN, x, y, x, y+n if piece_is_white else y-n)

def move_extract_dest(move):
    return move[MOVE_DEST_X], move[MOVE_DEST_Y]

def generate_regular_pawn_moves_forward(position, white_turn):
    return filter(
        lambda move: move_extract_dest(move) not in position,

        chain(
        (
            move_pawn_forward(x, y, white_turn)
            for x, y in pawns_of_turn(position, white_turn)
        ),
        (
            move_pawn_forward(x, y, white_turn, n=2)
            for x, y in pawns_of_turn(position, white_turn)
            if ( (y == PAWN_START and white_turn) or 
                 (y==BOARD_SIZE-PAWN_START+1 and not white_turn) )
            if (x, y+1 if white_turn else y-1) not in position
        ),
        ) # chain
    )

def generate_pawn_captures(position, white_turn):
    return ( (PAWN, x, y, x+d, y+1 if white_turn else y-1)
             for x, y in pawns_of_turn(position, white_turn)
             for d in (1, -1)
             if (x+d, y+1 if white_turn else y-1) in position
             if position[(x+d, y+1 if white_turn else y-1)][COLOUR_INDEX] 
             != white_turn
         )

def expand_pawn_moves_to_promotions(pawn_moves):
    for move in pawn_moves:
        if move[MOVE_DEST_Y] in (FIRST_RANK, LAST_RANK):
            for piece in (KNIGHT, BISHOP, ROOK, QUEEN):
                yield move + (piece, )
        else:
            yield move + (None,)

# yuck, this is broken and making changes to position 
def in_check_after_move(
        position, white_turn, orig_x, orig_y, dest_x, dest_y, 
        king_x=None, king_y=None):
    if None in (king_x, king_y):
        for cords, (piece_is_white, piece_name) in position.items():
            if white_turn == piece_is_white and piece_name == KING:
                king_x, king_y = cords
                break

    # save current piece
    orig_cords = (orig_x, orig_y)
    dest_cords = (dest_x, dest_y)
    orig_colour_and_piece = position[orig_cords]
    assert(orig_colour_and_piece[COLOUR_INDEX] == white_turn)
    orig_dest_colour_and_piece = (
        None if dest_cords not in position
        else position[dest_cords]
        )
    if None != orig_dest_colour_and_piece:
        assert(orig_dest_colour_and_piece[COLOUR_INDEX] != white_turn)
    assert(position[(king_x, king_y)][COLOUR_INDEX] == white_turn)

    del position[orig_cords]
    position[dest_cords] = orig_colour_and_piece
    
    result = any(
        king_x == move_dest_x and king_y == move_dest_y
        for (piece_name, 
             move_initial_x, move_initial_y, move_dest_x, move_dest_y, flags)
        in generate_moves_that_could_capture(position, not white_turn)
        )
    # restore position

    del position[dest_cords]
    if None != orig_dest_colour_and_piece:
        position[dest_cords] = orig_dest_colour_and_piece

    position[orig_cords] = orig_colour_and_piece 
    
    return result
            
def generate_castling(position, white_turn):
    y = FIRST_RANK if white_turn else LAST_RANK   
    # king needs to be on initial square
    if ( (KING_START_FILE, y) in position and 
         position[(KING_START_FILE, y)] == (KING, white_turn) ):
        return (
            (KING, KING_START_FILE, y, x+d, y, CASTLEING)
            for king_or_queenside, d, rook_x in 
            ( (KINGSIDE, 2, BOARD_SIZE), (QUEENSIDE, -2, 1) )
            if position[(white_turn, king_or_queenside)][COLOUR_INDEX]
            if empty_squares_in_between(
                    position, KING_START_FILE, y, rook_x, y)
            # check for being in check on the square in between
            # the destination square will be checked elsewhere with the 
            # routine that checks any move isn't going in to check
            if not in_check_after_move(position, white_turn, 
                                       x, y, x+neg_one_zero_or_pos_one(d), y, 
                                       x+neg_one_zero_or_pos_one(d), y,
                                   )
            )
    return ()
            
def generate_moves_that_could_capture(position, white_turn):
    return chain( generate_possible_piece_moves(position, white_turn),
                  expand_pawn_moves_to_promotions(
                      generate_pawn_captures(position, white_turn) )
    )
    
def generate_regular_piece_and_pawn_moves(position, white_turn):
    return chain( generate_moves_that_could_capture(position, white_turn),
                  expand_pawn_moves_to_promotions(
                    generate_regular_pawn_moves_forward(position, white_turn)),
                  generate_castling(position, white_turn)
    )

def generate_moves(position, white_turn):
        return filter( 
            lambda move: not in_check_after_move(
            position, white_turn, 
            move[MOVE_X], move[MOVE_Y], move[MOVE_DEST_X], move[MOVE_DEST_Y]),
            
            generate_regular_piece_and_pawn_moves(position, white_turn)
    )

def make_move(
        position, white_turn,
        piece_name, x, y, dest_x, dest_y, flags):
    assert( position[(x,y)] == (white_turn, piece_name) )
    del position[(x,y)]
    position[(dest_x, dest_y)] = (white_turn, piece_name)

test.py

#!/usr/bin/env python3

from move_generation import (
    PIECE_MOVES, generate_moves, make_move)
from initial_position import INITIAL_POSITION
from constants import KNIGHT, KING, BISHOP, ROOK, QUEEN, WHITE, BLACK

# check for correctness, don't allow duplicates
assert( all( len(move) == len(set(moves)) )
             for moves in PIECE_MOVES.values() )
# check for correct number of possible moves
assert( len(PIECE_MOVES[KNIGHT]) == 8 )
assert( len(PIECE_MOVES[KING]) == 8 )
assert( len(PIECE_MOVES[BISHOP]) == 28 )
assert( len(PIECE_MOVES[ROOK]) == 28 )
assert( len(PIECE_MOVES[QUEEN]) == 28+28 )


print( len(INITIAL_POSITION), len(INITIAL_POSITION.copy()) )

print( len(tuple(
     generate_moves(INITIAL_POSITION, WHITE) ) ) )

print( len(INITIAL_POSITION), len(INITIAL_POSITION.copy()) )

position = INITIAL_POSITION.copy()
turn = WHITE
i = 0
max_i = 2 ** 32
#while True:
    #moves  = tuple(generate_regular_piece_and_pawn_moves(position, turn))
    #move = moves[ i % len(moves) ]
    #make_
    #i = (i + 1) % max_i