from itertools import product from unicodedata import lookup from copy import deepcopy def cross(A, B): """Return sum of of every element from A and B.""" return tuple(a+b for a in A for b in B) pieces = ["pawn", "knight", "bishop", "rook", "queen", "king"] colors = ["white", "black"] ranks = "87654321" files = "abcdefgh" squares = cross(files, ranks) pawn_ranks = {"white" : '2', "black" : '7' } home_ranks = {"white" : '1', "black" : '8' } init_positions = {"pawn" : cross(files, pawn_ranks.values()), "rook" : cross("ah", home_ranks.values()), "knight" : cross("bg", home_ranks.values()), "bishop" : cross("cf", home_ranks.values()), "queen" : cross("d", home_ranks.values()), "king" : cross("e", home_ranks.values()) } AN_names = {'R' : "rook", 'N' : "knight", 'B' : "bishop", 'Q' : "queen", 'K' : "king"} def move(sq, v): """Add vector v to square and return a square, even if not in board.""" return chr(ord(sq[0]) + v[0]) + str(int(sq[1]) + v[1]) def moves(sq, dirs): """Add multiple vectors to a square. Return squares that are in board.""" targets = [move(sq, dir) for dir in dirs] return [sq for sq in targets if sq in squares] def invert(color): """Return inverted color.""" if color == "white": return "black" return "white" class Empty: """Empty class to avoid unnecessary quotations.""" def __repr__(self): return str(vars(self)) def parse_AN(move): """Parse chess move writen in algebraic notation. Return relevant information encapsulated within an Empty class, or False if parsing fails.""" def suffix_in(l, dval=False, i=1): nonlocal move x = move[-i:] if x in list(l): move = move[:-i] return x return dval queenside = ["{}-{}-{}".format(c,c,c) for c in "0O"] kingside = ["{}-{}".format(c,c) for c in "0O"] r = Empty() r.kingside = r.queenside = False if move in queenside: r.queenside = True return r if move in kingside: r.kingside = True return r r.move = move r.mate = suffix_in('#') == '#' r.check = suffix_in('+') == '+' r.promotion = suffix_in(AN_names, None) if r.promotion: r.promotion = AN_names[r.promotion] _ = suffix_in('=') r.dest = suffix_in(squares, i=2) r.capture = suffix_in("xX") != False r.rank = suffix_in(ranks, None) r.file = suffix_in(files, None) r.piece = suffix_in(AN_names, 'pawn') if len(move) > 0 or (r.promotion and r.piece != "pawn") or not r.dest: return False if r.piece != "pawn": r.piece = AN_names[r.piece] return r # Names for vectors up, down = (0, 1), (0, -1) left, right = (-1, 0), (1, 0) rook_dirs = (up, down, left, right) upper_left, upper_right = (-1, 1), (1, 1) lower_left, lower_right = (-1, -1), (1, -1) bishop_dirs = (upper_left, upper_right, lower_left, lower_right) class Piece: """Contains piece color and piece type""" def __init__(self, color=None, piece=None): """Assign color and piece. If not specified they default to None.""" self.color = color self.piece = piece def __repr__(self): """Returns unicode representation of instance.""" if None in (self.piece, self.color): return "" name = self.color.upper() + " CHESS " + self.piece.upper() return lookup(name) def __str__(self): """Same as repr.""" return repr(self) def __eq__(self, other): """Ensures equality testing. If color or piece or both equal None, instance equals None.""" if other == None: return None in (self.color, self.piece) if not isinstance(other, Piece): return NotImplemented return self.piece == other.piece and self.color == other.color class Game: """Chess game class. Squares are represented by strings made of letters and integers e.g. 'h4'. Board is represented by a dictionary whose keys are squares and values instances of Piece class. Moves are represented internally by tuples of squares; a source and destination. Moves are appended to moves stack. Upon each move, the relevant attributes of the instance (the game state) are appended to the history stack. This allows loading previous game states to instance variables.""" def __init__(self): self.index = 0 self.make_board() self.moves = [] self.AN_moves = [] self.history = [] self.turn = "white" self.can_castle = {"white" : {"queen" : True, "king" : True}, "black" : {"queen" : True, "king" : True} } self.history.append((deepcopy(self.can_castle), deepcopy(self.board), self.turn)) def __repr__(self): """Return unicode board representation.""" r = "" for rank in ranks: r += rank + " |" for file in files: sq = file+rank r += " " if self.is_empty(sq): r += " " else: r += repr(self.board[sq]) r += "\n" r += " +" + "-"*16 + "\n" r += " "*4 + " ".join(list(files)) return r def make_board(self): """Create board and store it in self.board.""" self.board = dict((sq, Piece()) for sq in squares) # Add pieces for piece, positions in init_positions.items(): for sq in positions: self.board[sq].piece = piece # Add colors for sq in cross(files, "12"): self.board[sq].color = "white" for sq in cross(files, "78"): self.board[sq].color = "black" # Castling markers self.board["a1"].side = "queen" self.board["a8"].side = "queen" self.board["h1"].side = "king" self.board["h8"].side = "king" def rebase(self): """Shorten moves and history stacks if move is made while past game state is loaded into memory. Essentialy overrides future game states.""" self.moves = self.moves[:self.index] self.history = self.history[:self.index+1] def timetravel(self, i): """Load game state from past into instance variables. Can be reversed since history and moves stack stay intact.""" if i >= len(self.history) or i < 0: return False c, board, turn = self.history[i] self.board = board self.turn = turn self.can_castle = c self.index = i def prev(self): """Load game state from previous turn, if it exists.""" return self.timetravel(self.index-1) def next(self): """Load game state from next turn, if it exists.""" return self.timetravel(self.index+1) def last(self): """Load game state from last move.""" return self.timetravel(len(self.history)-1) def first(self): """Load first recorded game state, an empty board.""" return self.timetravel(0) def is_empty(self, sq): """Return True if square is empty, else False.""" return self.board[sq] == None def is_check(self, color=None): """Return True if king of active player is in check else False.""" if color == None: color = self.turn sq = self.occupying(Piece(color, "king"))[0] return self.is_attacked(color, sq) def is_mate(self): """Return True if checkmate for active player.""" return self.is_check() and self.is_stall() def is_stall(self): """Return True if active player has no legal move. Whether the king is in check is irrelevant.""" for move in self.all_legal_moves(): new = deepcopy(self) if self.board[move[0]].piece == "pawn" and move[1][1]==home_ranks[invert(self.turn)]: move += ("queen",) if new.move(*move): return False return True def endofgame(self): """Return True if active player can't make a legal move.""" return self.is_stall() or self.is_mate() def occupying(self, piece): """Return squares occupied by piece. Piece variable is an instance of Piece class.""" return [sq for sq in squares if self.board[sq] == piece] def castle(self, side): """Castle if possible, if not return False.""" color = self.turn rank = home_ranks[self.turn] if side not in ("queen", "king"): return False if not self.can_castle[color][side]: return False if side == "queen": check_if_empty = cross("bcd", rank) check_if_attacked = cross("cde", rank) rook_move = "a"+rank, "d"+rank king_move = "e"+rank, "c"+rank elif side == "king": check_if_empty = cross("fg", rank) check_if_attacked = check_if_empty rook_move = "h"+rank, "f"+rank king_move = "e"+rank, "g"+rank are_empty = all(map(self.is_empty, check_if_empty)) not_attacked = not any(self.is_attacked(color, sq) for sq in check_if_attacked) if not (are_empty and not_attacked): return False self.rebase() self.board[king_move[1]] = self.board[king_move[0]] self.board[rook_move[1]] = self.board[rook_move[0]] self.board[king_move[0]] = Piece() self.board[rook_move[0]] = Piece() self.can_castle[color]["queen"] = False self.can_castle[color]["king"] = False self.moves.append(king_move) self.turn = invert(self.turn) self.history.append((deepcopy(self.can_castle), deepcopy(self.board), self.turn)) self.index += 1 return True def move(self, source, target, promotion=None): """Move piece from source to target if legal and promote it to promotion if it's a pawn whos reached the last rank. Return False if fails.""" new = deepcopy(self) if not new.is_legal(source, target, promotion): return False board = new.board moved = board[source] eaten = board[target] if promotion == None: promotion = moved else: promotion = Piece(new.turn, promotion) if moved.piece == "pawn" and target in moves(source, bishop_dirs) and eaten==None: # En passant a, b = move(target, up), move(target, down) board[a] = Piece() board[b] = Piece() # Castling if moved.piece == "rook": new.can_castle[new.turn][moved.side] = False if eaten.piece == "rook": new.can_castle[eaten.color][eaten.side] = False if moved.piece == "king": new.can_castle[new.turn]["queen"] = False new.can_castle[new.turn]["king"] = False board[source] = Piece() board[target] = promotion if new.is_check(self.turn): return False self.rebase() self.moves.append((source, target)) self.board = new.board self.can_castle = new.can_castle self.turn = invert(moved.color) self.history.append((deepcopy(self.can_castle), deepcopy(self.board), self.turn)) self.index += 1 return True def is_legal(self, source, target, promotion): """Return True if move is partially legal. Does not check whether the active players king is in check after the move! See the validate_move method.""" piece = self.board[source].piece color = self.board[source].color pawn_on_last_rank = piece == "pawn" and \ home_ranks[invert(color)] == target[1] if promotion != None and not pawn_on_last_rank: return False if promotion == None and pawn_on_last_rank: return False if promotion in ("king", "pawn"): return False return color == self.turn and target in self.possible_moves(source) def is_attacked(self, color, sq): """Return True if square is attacked by enemy of opposite color.""" enemy_color = invert(color) attacked = [] for piece in pieces: occupied = self.occupying(Piece(enemy_color, piece)) attacked += sum(list(map(self.attacks, occupied)), []) return sq in attacked def all_legal_moves(self): """Return all possible partially legal moves by active player. These moves can put the active player in check, thus they are not completely legal. Castlings not included.""" color = self.turn sqrs = [] for piece in pieces: piece = Piece(color, piece) sqrs += self.occupying(piece) r = [] for sq in sqrs: mvs = self.possible_moves(sq) r += [(sq, m) for m in mvs] return r def possible_moves(self, sq): """Return all squares the piece from sq can move to. Moving to these squares might put the player in check. Castlings not included.""" board = self.board piece = board[sq].piece color = board[sq].color def can_eat(target): eaten = board[target] return eaten != None and eaten.color != color is_empty = self.is_empty targets = [t for t in self.attacks(sq) if can_eat(t) or is_empty(t)] if piece == "pawn": r = [] dir, back = up, down if color == "black": dir, back = down, up frwd = move(sq, dir) jump = move(frwd, dir) if frwd in squares and is_empty(frwd): r.append(frwd) is_on_pawn_rank = pawn_ranks[color] == sq[1] if is_on_pawn_rank and is_empty(jump): r.append(jump) for t in targets: a, b = move(t, dir), move(t, back) en_passant = can_eat(b) and self.index>0 and\ self.moves[self.index-1] == (a,b) if can_eat(t) or en_passant: r.append(t) return r elif piece == "king": return [sq for sq in targets if not self.is_attacked(color, sq)] else: return targets def attacks(self, sq): """Return all squares the piece on sq attacks. Return empty list if square is empty.""" if sq not in squares: return [] board = self.board piece = board[sq].piece color = board[sq].color def possible_line(dir): target = move(sq, dir) while target in board and self.is_empty(target): yield target target = move(target, dir) if target in board: yield target def possible_lines(dirs): return sum(list(map(list, map(possible_line, dirs))), []) rook = possible_lines(rook_dirs) bishop = possible_lines(bishop_dirs) if piece == "pawn": if color == "white": dirs = upper_left, upper_right else: dirs = lower_left, lower_right return moves(sq, dirs) elif piece == "knight": nums = [2, -2, 1, -1] dirs = [(a, b) for a,b in product(nums, nums) if abs(a) != abs(b)] return moves(sq, dirs) elif piece == "rook": return rook elif piece == "bishop": return bishop elif piece == "queen": return rook + bishop elif piece == "king": dirs = rook_dirs + bishop_dirs return moves(sq, dirs) return [] def validate_move(self, source, dest, promotion): """Create a copy of the game and try given move. Return True if the move succeeds. This method exists because is_legal does not test whether the move puts the player in check.""" new = deepcopy(self) valid = new.move(source, dest, promotion) new.turn = self.turn return valid, new def AN_move(self, move): """Move is given in algebraic notation. The method moves the piece if the move is legal and unambiguous, else it returns False.""" AN = move info = parse_AN(move) if info == False: return False if info.queenside: return self.castle("queen") if info.kingside: return self.castle("king") piece = info.piece sqrs = self.occupying(Piece(self.turn, piece)) promotion = info.promotion dest = info.dest if info.file != None: sqrs = [sq for sq in sqrs if sq[0]==info.file] if info.rank != None: sqrs = [sq for sq in sqrs if sq[1]==info.rank] moves = [] for sq in sqrs: valid, new = self.validate_move(sq, dest, promotion) check = new.is_check()==info.check and new.is_mate()==info.mate if valid and check: moves.append((sq, dest, promotion)) if len(moves) != 1: return False move = moves[0] self.move(*move) self.AN_moves.append(AN) return True game = Game() while True: i = input(">> ") game.AN_move(i) if i =="prev": game.prev() print(game) print(game.index, len(game.history), len(game.moves))