diff --git a/chess.py b/chess.py index 639b0a2..13f4296 100644 --- a/chess.py +++ b/chess.py @@ -1,3 +1,4 @@ +from itertools import product from unicodedata import lookup def cross(A, B): @@ -17,6 +18,14 @@ init_positions = {"pawn" : cross(files, pawn_ranks), "king" : cross("e", home_ranks) } +def addv(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] + +def negatevs(vectors): + return [(-a, -b) for a,b in vectors] class Piece: def __init__(self, color=None, piece=None): @@ -75,6 +84,47 @@ class Game: def occupying(self, piece): return [sq for sq in squares if self.board[sq] == piece] + 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), []) + + piece = self.board[sq].piece + + if piece == "pawn": + pass + 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)] + elif piece == "rook": + return possible_rook() + elif piece == "bishop": + return possible_bishop() + elif piece == "queen": + return possible_bishop() + possible_rook() @@ -82,7 +132,8 @@ def test(): game = Game() assert len(squares) == 8**2 assert sum(map(len, init_positions.values())) == 8*4 + game.board["d4"] = Piece("black", "bishop") print(game) - print(game.occupying(Piece("white", "pawn"))) + print (game.possible_moves("d4")) test()