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 = "27" home_ranks = "18" init_positions = {"pawn" : cross(files, pawn_ranks), "rook" : cross("ah", home_ranks), "knight" : cross("bg", home_ranks), "bishop" : cross("cf", home_ranks), "queen" : cross("d", home_ranks), "king" : cross("e", home_ranks) } def get_rank(sq): return int(sq[1]) def get_file(sq): return sq[0] 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" 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 = [] 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): if not self.is_legal(source, target): return False board = self.board moved = board[source] eaten = board[target] self.moves.append((source, target)) self.stack.append(eaten) board[source] = Piece() board[target] = moved return True def is_legal(self, source, target): return 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 = up if color == "black": dir = down frwd = move(sq, dir) jump = move(frwd, dir) if frwd in squares and is_empty(frwd): r.append(frwd) is_on_pawn_rank = color == "white" and get_rank(sq)==2 \ or \ color == "black" and get_rank(sq)==7 if is_on_pawn_rank and is_empty(jump): r.append(jump) return r + [sq for sq in targets if can_eat(sq)] 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", "a3"), ("b1", "c3"), ("c3", "b5"), ("b5", "c7")] 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()