diff --git a/chess.py b/chess.py index 452375e..9c197e4 100644 --- a/chess.py +++ b/chess.py @@ -4,6 +4,9 @@ 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) @@ -18,14 +21,17 @@ init_positions = {"pawn" : cross(files, pawn_ranks), "king" : cross("e", home_ranks) } -def addv(sq, v): +def move(sq, v): return chr(ord(sq[0]) + v[0]) + str(int(sq[1]) + v[1]) -def addvs(sq, vectors): - return [addv(sq, v) for v in vectors] +up, down = (0, 1), (0, -1) +left, right = (-1, 0), (1, 0) +rook_dirs = (up, down, left, right) -def negatevs(vectors): - return [(-a, -b) for a,b in vectors] +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) class Piece: def __init__(self, color=None, piece=None): @@ -107,77 +113,73 @@ class Game: return False def possible_moves(self, sq): - def is_possible(target): - return target in squares and \ - self.board[target].color != self.board[sq].color - def are_possible(vecs): - r = [] - sqrs = addvs(sq, vecs) - for s in sqrs: - if not is_possible(s): - break - r.append(s) - if not self.is_empty(s): - break - return r - - def possible_rook(): - rvecs = list(zip(range(1,8), [0]*7)) - fvecs = [(b,a) for a,b in rvecs] - vecs = [rvecs, fvecs, negatevs(rvecs), negatevs(fvecs)] - return sum(map(are_possible, vecs), []) - - def possible_bishop(): - ldvecs = list(zip(range(1,8), range(1,8))) - rdvecs = list(zip(range(1,8), range(-1, -8, -1))) - vecs = [ldvecs, rdvecs, negatevs(ldvecs), negatevs(rdvecs)] - return sum(map(are_possible, vecs), []) + if sq not in squares: + return [] board = self.board 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): + yield target + target = move(target, dir) + if can_eat(target): + 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": r = [] - direction = (0,1) + dir = up if color == "black": - direction = (0, -1) - frwd = addv(sq, direction) - jump = addv(frwd, direction) - eatl = addv(frwd, (-1, 0)) - eatr = addv(frwd, (1, 0)) + dir = down + frwd = move(sq, dir) + jump = move(frwd, dir) + eat = (move(frwd, left), + move(frwd, right)) - if is_possible(frwd) and self.is_empty(frwd): + if can_move_to(frwd): r.append(frwd) - checklr = lambda x: is_possible(x) \ - and \ - not self.is_empty(x) \ - and \ - board[x].color != color - if checklr(eatl): - r.append(eatl) - if checklr(eatr): - r.append(eatr) - is_on_pawn_rank = (color == "black" and sq[1]=='7') \ - or \ - (color == "white" and sq[1]=='2') - if is_possible(jump) and frwd in r and self.is_empty(jump) and is_on_pawn_rank: - r.append(jump) - return r + 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)] + elif piece == "knight": - vecs = list(product((1,-1), (2,-2))) - vecs += list(product((2,-2), (1,-1))) - return [sq for sq in addvs(sq, vecs) if is_possible(sq)] + 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)] elif piece == "rook": - return possible_rook() + return rook elif piece == "bishop": - return possible_bishop() + return bishop elif piece == "queen": - return possible_bishop() + possible_rook() + return rook + bishop elif piece == "king": - vecs = product([-1, 0, 1], [-1, 0, 1]) - return [sq for sq in addvs(sq, vecs) - if is_possible(sq) and not self.is_attacked(sq)] + 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)] return []