from itertools import product from unicodedata import lookup def cross(A, 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): return chr(ord(sq[0]) + v[0]) + str(int(sq[1]) + v[1]) def moves(sq, dirs): targets = [move(sq, dir) for dir in dirs] return [sq for sq in targets if sq in squares] def invert(color): if color == "white": return "black" return "white" class Empty: def __repr__(self): return str(vars(self)) def parse_AN(move): 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) 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 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: def __init__(self, color=None, piece=None): self.color = color self.piece = piece def __repr__(self): if None in (self.piece, self.color): return None name = self.color.upper() + " CHESS " + self.piece.upper() return lookup(name) def __eq__(self, other): 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: def __init__(self): self.make_board() self.moves = [] self.stack = [] self.turn = "white" def __repr__(self): # 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): 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" def is_empty(self, sq): return self.board[sq] == None def occupying(self, piece): return [sq for sq in squares if self.board[sq] == piece] def move(self, source, target, promotion=None): if not self.is_legal(source, target, promotion): return False board = self.board moved = board[source] eaten = board[target] if promotion == None: promotion = moved else: promotion = Piece(self.turn, promotion) if moved.piece == "pawn" and target in moves(source, bishop_dirs): # En passant a, b = move(target, up), move(target, down) board[a] = Piece() board[b] = Piece() self.moves.append((source, target)) self.stack.append(eaten) board[source] = Piece() board[target] = promotion self.turn = invert(moved.color) return True def is_legal(self, source, target, promotion): 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 == "pawn": return False return color == self.turn and target in self.possible_moves(source) def is_attacked(self, color, sq): 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 possible_moves(self, sq): 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.moves[-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): 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 test(): game = Game() assert len(squares) == 8**2 assert sum(map(len, init_positions.values())) == 8*4 moves = [("a2", "a4"), ("b8", "c6"), ("a4", "a5"), ("b7", "b5"), ("a5", "b6"), ("c6", "e5"), ("b6", "b7"), ("h7", "h5"), ("b7", "b8", "pawn")] for m in moves: game.move(*m) print (game) game.board["c3"] = Piece("black", "king") print (game.attacks("c3")) print(game.possible_moves("c3")) print(game.is_attacked("black", "d8")) test()