Refactored possible_moves method

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):
+up,   down    = (0, 1),  (0, -1)
-    return [addv(sq, v) for v in vectors]
+left, right   = (-1, 0), (1, 0)
+rook_dirs     = (up, down, left, right)
 
-def negatevs(vectors):
+up_left,   up_right   = (-1, 1),  (1, 1)
-    return [(-a, -b) for a,b in vectors]
+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):
+        if sq not in squares:
-            return target in squares and \
+            return []
-                   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), [])
 
         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)
+                dir = down
-            frwd  = addv(sq, direction)
+            frwd = move(sq,   dir)
-            jump = addv(frwd, direction)
+            jump = move(frwd, dir)
-            eatl = addv(frwd, (-1, 0))
+            eat = (move(frwd, left),
-            eatr = addv(frwd,  (1, 0))
+                   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) \
+                is_on_pawn_rank = (color == "black" and sq[1]=='7') \
-                                and \
+                                   or \
-                                not self.is_empty(x) \
+                                  (color == "white" and sq[1]=='2')
-                                and \
+                if can_move_to(jump) and is_on_pawn_rank:
-                                board[x].color != color
+                    r.append(jump)
-            if checklr(eatl):
+            return r + [sq for sq in eat if can_eat(sq)]
-                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
         elif piece == "knight":
-            vecs =  list(product((1,-1), (2,-2)))
+            nums = [2, -2, 1, -1]
-            vecs += list(product((2,-2), (1,-1)))
+            vectors = [(a, b) for a,b in product(nums, nums)
-            return [sq for sq in addvs(sq, vecs) if is_possible(sq)]
+                                  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])
+            vectors = rook_dirs + bishop_dirs
-            return [sq for sq in addvs(sq, vecs)
+            targets = [move[sq, v] for v in vectors]
-                       if is_possible(sq) and not self.is_attacked(sq)]
+            return [sq for sq in targets
+                           if can_move_to(sq) and not is_attacked(sq)]
         return []