From fbf001a9b2ff138a01f4743822e3b3d86f6a137e Mon Sep 17 00:00:00 2001
From: Tibor Bizjak <tibor.bizjak@gmail.com>
Date: Tue, 10 Sep 2019 15:22:15 +0200
Subject: [PATCH] Added logic so move is invalid if it opens a check

---
 chess.py | 60 ++++++++++++++++++++++++++++++--------------------------
 1 file changed, 32 insertions(+), 28 deletions(-)

diff --git a/chess.py b/chess.py
index 24aa9dc..8f60d6e 100644
--- a/chess.py
+++ b/chess.py
@@ -170,13 +170,14 @@ class Game:
         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
+        return self.is_check() and self.is_stall()
+
+    def is_stall(self):
         for move in self.all_legal_moves():
             new = deepcopy(self)
-            new.move(*move)
-            if not new.is_check(self.turn):
+            if self.board[move[0]].piece == "pawn" and move[1][1]==home_ranks[invert(self.turn)]:
+                move += ("queen",)
+            if new.move(*move):
                 return False
         return True
 
@@ -215,18 +216,22 @@ class Game:
         self.board[rook_move[0]] = Piece()
         self.can_castle[color]["queen"] = False
         self.can_castle[color]["king"]  = False
+        self.moves.append(king_move)
+        return True
 
     def move(self, source, target, promotion=None):
-        if not self.is_legal(source, target, promotion):
+        new = deepcopy(self)
+        if not new.is_legal(source, target, promotion):
             return False
-        board = self.board
+
+        board = new.board
         moved = board[source]
         eaten = board[target]
 
         if promotion == None:
             promotion = moved
         else:
-            promotion = Piece(self.turn, promotion)
+            promotion = Piece(new.turn, promotion)
 
         if moved.piece == "pawn" and target in moves(source, bishop_dirs) and eaten==None:
             # En passant
@@ -236,18 +241,22 @@ class Game:
 
         # Castling
         if moved.piece == "rook":
-            self.can_castle[self.turn][moved.side] = False
+            new.can_castle[new.turn][moved.side] = False
         if eaten.piece == "rook":
-            self.can_castle[eaten.color][eaten.side] = False
+            new.can_castle[eaten.color][eaten.side] = False
         if moved.piece == "king":
-            self.can_castle[self.turn]["queen"] = False
-            self.can_castle[self.turn]["king"]  = False
-
-        self.moves.append((source, target))
-        self.stack.append(eaten)
+            new.can_castle[new.turn]["queen"] = False
+            new.can_castle[new.turn]["king"]  = False
 
         board[source] = Piece()
         board[target] = promotion
+
+        if new.is_check(self.turn):
+            return False
+
+        self.moves.append((source, target))
+        self.board = new.board
+        self.can_castle = new.can_castle
         self.turn = invert(moved.color)
         return True
 
@@ -370,18 +379,13 @@ class Game:
 
 def test():
     game = Game()
-    assert len(squares) == 8**2
-    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'),
-             ('e2', 'e3'), ('d6', 'f4'), ('e1', 'd1'), ('f4', 'g4')]
-    for m in moves:
-        game.move(*m)
-        print (game.is_check())
-        print (game)
-
-    game.board['e1'] = game.board['c1'] = game.board['d2'] = Piece("white","pawn")
-    print (game.is_mate(), "mate")
-    print(game)
+    game.board = {sq : Piece() for sq in squares}
+    game.board['a8'] = Piece("white", "king")
+    game.board['f6'] = Piece("black", "pawn")
+    game.board['f5'] = Piece("white", "pawn")
+    game.board['c7'] = Piece("black", "queen")
+    game.board['h1'] = Piece("black", "king")
+    print(game, game.is_stall())
+    print(game.all_legal_moves())
 
 test()