From 582a635e17c7a91f23753f69713169c8247d0194 Mon Sep 17 00:00:00 2001
From: Tibor Bizjak <tibor.bizjak@gmail.com>
Date: Tue, 10 Sep 2019 14:30:43 +0200
Subject: [PATCH] Added is_check and is_mate methods and helper method
 all_legal_moves

---
 chess.py | 41 ++++++++++++++++++++++++++++++++++-------
 1 file changed, 34 insertions(+), 7 deletions(-)

diff --git a/chess.py b/chess.py
index 3c28193..24aa9dc 100644
--- a/chess.py
+++ b/chess.py
@@ -1,5 +1,6 @@
 from itertools import product
 from unicodedata import lookup
+from copy import deepcopy
 
 def cross(A, B):
     return tuple(a+b for a in A for b in B)
@@ -162,6 +163,23 @@ class Game:
     def is_empty(self, sq):
         return self.board[sq] == None
 
+    def is_check(self, color=None):
+        if color == None:
+            color = self.turn
+        sq = self.occupying(Piece(color, "king"))[0]
+        return self.is_attacked(color, sq)
+
+    def is_mate(self):
+        sq = self.occupying(Piece(self.turn, "king"))[0]
+        if not self.is_check():
+            return False
+        for move in self.all_legal_moves():
+            new = deepcopy(self)
+            new.move(*move)
+            if not new.is_check(self.turn):
+                return False
+        return True
+
     def occupying(self, piece):
         return [sq for sq in squares if self.board[sq] == piece]
 
@@ -254,6 +272,18 @@ class Game:
             attacked += sum(list(map(self.attacks, occupied)), [])
         return sq in attacked
 
+    def all_legal_moves(self):
+        color = self.turn
+        sqrs = []
+        for piece in pieces:
+            piece = Piece(color, piece)
+            sqrs += self.occupying(piece)
+        r = []
+        for sq in sqrs:
+            mvs = self.possible_moves(sq)
+            r += [(sq, m) for m in mvs]
+        return r
+
     def possible_moves(self, sq):
         board = self.board
         piece = board[sq].piece
@@ -344,17 +374,14 @@ def test():
     assert sum(map(len, init_positions.values())) == 8*4
     moves = [('d2', 'd4'), ('d7', 'd5'), ('d1', 'd3'), ('h7', 'h5'),
              ('c1', 'h6'), ('g7', 'h6'), ('b1', 'c3'), ('d8', 'd6'),
-             ('h2', 'h4'), ('d6', 'f4'), ('e2', 'e3'), ('f4', 'g4'), ('e1', 'd1')]
+             ('e2', 'e3'), ('d6', 'f4'), ('e1', 'd1'), ('f4', 'g4')]
     for m in moves:
         game.move(*m)
-        print (game.can_castle)
+        print (game.is_check())
         print (game)
 
-    game.castle("queen")
+    game.board['e1'] = game.board['c1'] = game.board['d2'] = Piece("white","pawn")
+    print (game.is_mate(), "mate")
     print(game)
-    print (game.attacks("c3"))
-    print(game.possible_moves("c3"))
-    print(game.is_attacked("black", "d8"))
-
 
 test()