diff --git a/chess.py b/chess.py index 9c197e4..9a817f7 100644 --- a/chess.py +++ b/chess.py @@ -21,17 +21,32 @@ init_positions = {"pawn" : cross(files, pawn_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) -up_left, up_right = (-1, 1), (1, 1) -down_left, down_right = (-1, -1), (1, -1) -bishop_dirs = (up_left, up_right, - down_left, down_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): @@ -109,10 +124,50 @@ class Game: def is_legal(self, source, target): return target in self.possible_moves(source) - def is_attacked(self, sq): - return False + 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 [] @@ -120,21 +175,12 @@ class Game: piece = board[sq].piece color = board[sq].color - def can_eat(target): - if target not in board: - return False - eaten = board[target] - return eaten != None and color != eaten.color - - def can_move_to(sq): - return sq in board and self.is_empty(sq) - def possible_line(dir): target = move(sq, dir) - while can_move_to(target): + while target in board and self.is_empty(target): yield target target = move(target, dir) - if can_eat(target): + if target in board: yield target def possible_lines(dirs): @@ -144,31 +190,16 @@ class Game: bishop = possible_lines(bishop_dirs) if piece == "pawn": - r = [] - dir = up - if color == "black": - dir = down - frwd = move(sq, dir) - jump = move(frwd, dir) - eat = (move(frwd, left), - move(frwd, right)) - - if can_move_to(frwd): - r.append(frwd) - is_on_pawn_rank = (color == "black" and sq[1]=='7') \ - or \ - (color == "white" and sq[1]=='2') - if can_move_to(jump) and is_on_pawn_rank: - r.append(jump) - return r + [sq for sq in eat if can_eat(sq)] - + 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] - vectors = [(a, b) for a,b in product(nums, nums) - if abs(a) != abs(b)] - targets = [move(sq, v) for v in vectors] - return [sq for sq in targets - if can_move_to(sq) or can_eat(sq)] + 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": @@ -176,10 +207,8 @@ class Game: elif piece == "queen": return rook + bishop elif piece == "king": - vectors = rook_dirs + bishop_dirs - targets = [move[sq, v] for v in vectors] - return [sq for sq in targets - if can_move_to(sq) and not is_attacked(sq)] + dirs = rook_dirs + bishop_dirs + return moves(sq, dirs) return [] @@ -188,9 +217,15 @@ def test(): game = Game() assert len(squares) == 8**2 assert sum(map(len, init_positions.values())) == 8*4 - moves = [("a2", "a3"), ("b1", "c3"), ("d1", "d2")] + 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()