From c0f2df0b59637a1d2f77355f2334a8fc29d92eaf Mon Sep 17 00:00:00 2001
From: Tibor Bizjak <tibor.bizjak@gmail.com>
Date: Tue, 10 Sep 2019 13:08:06 +0200
Subject: [PATCH] Added en passant and pawn promotion logic

---
 chess.py | 30 +++++++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/chess.py b/chess.py
index 8380936..1c3b235 100644
--- a/chess.py
+++ b/chess.py
@@ -155,23 +155,42 @@ class Game:
     def occupying(self, piece):
         return [sq for sq in squares if self.board[sq] == piece]
 
-    def move(self, source, target):
-        if not self.is_legal(source, target):
+    def move(self, source, target, promotion=None):
+        if not self.is_legal(source, target, promotion):
             return False
         board = self.board
         moved = board[source]
         eaten = board[target]
+        if promotion == None:
+            promotion = moved
+        else:
+            promotion = Piece(self.turn, promotion)
+
+        if moved.piece == "pawn" and target in moves(source, bishop_dirs):
+            # En passant
+            a, b = move(target, up), move(target, down)
+            board[a] = Piece()
+            board[b] = Piece()
 
         self.moves.append((source, target))
         self.stack.append(eaten)
 
         board[source] = Piece()
-        board[target] = moved
+        board[target] = promotion
         self.turn = invert(moved.color)
         return True
 
-    def is_legal(self, source, target):
+    def is_legal(self, source, target, promotion):
+        piece = self.board[source].piece
         color = self.board[source].color
+        pawn_on_last_rank = piece == "pawn" and \
+                            home_ranks[invert(color)] == target[1]
+        if promotion != None and not pawn_on_last_rank:
+            return False
+        if promotion == None and pawn_on_last_rank:
+            return False
+        if promotion == "pawn":
+            return False
         return color == self.turn and target in self.possible_moves(source)
 
     def is_attacked(self, color, sq):
@@ -270,7 +289,8 @@ def test():
     game = Game()
     assert len(squares) == 8**2
     assert sum(map(len, init_positions.values())) == 8*4
-    moves = [("a2", "a4"), ("b8", "c6"), ("a4", "a5"), ("b7", "b5"), ("a5", "b6")]
+    moves = [("a2", "a4"), ("b8", "c6"), ("a4", "a5"), ("b7", "b5"), ("a5", "b6"),
+             ("c6", "e5"), ("b6", "b7"), ("h7", "h5"), ("b7", "b8", "pawn")]
     for m in moves:
         game.move(*m)
         print (game)