Further refactoring of possible_moves and helper functions

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)
+upper_left, upper_right = (-1, 1),  (1, 1)
-down_left, down_right = (-1, -1), (1, -1)
+lower_left, lower_right = (-1, -1), (1, -1)
-bishop_dirs           = (up_left,   up_right,
+bishop_dirs             = (upper_left, upper_right,
-                         down_left, down_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):
+    def is_attacked(self, color, sq):
-        return False
+        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 = []
+            if color == "white":
-            dir = up
+                dirs = upper_left, upper_right
-            if color == "black":
+            else:
-                dir = down
+                dirs = lower_left, lower_right
-            frwd = move(sq,   dir)
+            return moves(sq, dirs)
-            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)]
-
         elif piece == "knight":
             nums = [2, -2, 1, -1]
-            vectors = [(a, b) for a,b in product(nums, nums)
+            dirs = [(a, b) for a,b in product(nums, nums)
                                if abs(a) != abs(b)]
-            targets = [move(sq, v) for v in vectors]
+            return moves(sq, dirs)
-            return [sq for sq in targets
-                           if can_move_to(sq) or can_eat(sq)]
         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
+            dirs = rook_dirs + bishop_dirs
-            targets = [move[sq, v] for v in vectors]
+            return moves(sq, dirs)
-            return [sq for sq in targets
-                           if can_move_to(sq) and not is_attacked(sq)]
         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()