Single_Player_Chess/chess.py

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)

pieces = ["pawn", "knight", "bishop", "rook", "queen", "king"]
colors = ["white", "black"]

ranks   = "87654321"
files   = "abcdefgh"
squares = cross(files, ranks)

pawn_ranks = {"white" : '2',
              "black" : '7'
             }
home_ranks = {"white" : '1',
              "black" : '8'
             }
init_positions = {"pawn"   : cross(files, pawn_ranks.values()),
                  "rook"   : cross("ah",  home_ranks.values()),
                  "knight" : cross("bg",  home_ranks.values()),
                  "bishop" : cross("cf",  home_ranks.values()),
                  "queen"  : cross("d",   home_ranks.values()),
                  "king"   : cross("e",   home_ranks.values())
                 }

AN_names = {'R' : "rook",
            'N' : "knight",
            'B' : "bishop",
            'Q' : "queen",
            'K' : "king"}

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"

class Empty:
    def __repr__(self):
        return str(vars(self))

def parse_AN(move):
    def suffix_in(l, dval=False, i=1):
        nonlocal move
        x = move[-i:]
        if x in list(l):
            move = move[:-i]
            return x
        return dval

    queenside = ["{}-{}-{}".format(c,c,c) for c in "0O"]
    kingside  = ["{}-{}".format(c,c)      for c in "0O"]

    r = Empty()
    r.kingside = r.queenside = False
    if move in queenside:
        r.queenside = True
        return r
    if move in kingside:
        r.kingside = True
        return r

    r.move      = move
    r.mate      = suffix_in('#') == '#'
    r.check     = suffix_in('+') == '+'
    r.promotion = suffix_in(AN_names)
    if r.promotion:
        r.promotion = AN_names[r.promotion]
        _           = suffix_in('=')
    r.dest    = suffix_in(squares, i=2)
    r.capture = suffix_in("xX") != False
    r.rank    = suffix_in(ranks, None)
    r.file    = suffix_in(files, None)
    r.piece   = suffix_in(AN_names, 'pawn')
    if len(move) > 0 or (r.promotion and r.piece != "pawn") or not r.dest:
        return False
    if r.piece != "pawn":
        r.piece = AN_names[r.piece]
    return r

up,   down    = (0, 1),  (0, -1)
left, right   = (-1, 0), (1, 0)
rook_dirs     = (up, down, left, right)

upper_left, upper_right = (-1, 1),  (1, 1)
lower_left, lower_right = (-1, -1), (1, -1)
bishop_dirs             = (upper_left, upper_right,
                           lower_left, lower_right)

class Piece:
    def __init__(self, color=None, piece=None):
        self.color = color
        self.piece = piece

    def __repr__(self):
        if None in (self.piece, self.color):
            return None
        name = self.color.upper() + " CHESS " + self.piece.upper()
        return lookup(name)

    def __eq__(self, other):
        if other == None:
            return  None in (self.color, self.piece)
        if not isinstance(other, Piece):
            return NotImplemented
        return self.piece == other.piece and self.color == other.color

class Game:
    def __init__(self):
        self.make_board()
        self.moves = []
        self.stack = []
        self.turn  = "white"
        self.can_castle = {"white" : {"queen" : True,
                                      "king"  : True},
                           "black" : {"queen" : True,
                                      "king"  : True}
                          }

    def __repr__(self):
        # Unicode board representation
        r = ""
        for rank in ranks:
            r += rank + " |"
            for file in files:
                sq = file+rank
                r += " "
                if self.is_empty(sq):
                    r += " "
                else:
                    r += repr(self.board[sq])
            r += "\n"
        r += "  +" + "-"*16 + "\n"
        r += " "*4 + " ".join(list(files))
        return r

    def make_board(self):
        self.board = dict((sq, Piece()) for sq in squares)
        # Add pieces
        for piece, positions in init_positions.items():
            for sq in positions:
                self.board[sq].piece = piece
        # Add colors
        for sq in cross(files, "12"):
            self.board[sq].color = "white"
        for sq in cross(files, "78"):
            self.board[sq].color = "black"
        # Castling markers
        self.board["a1"].side = "queen"
        self.board["a8"].side = "queen"
        self.board["h1"].side = "king"
        self.board["h8"].side = "king"

    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):
        return self.is_check() and self.is_stall()

    def is_stall(self):
        for move in self.all_legal_moves():
            new = deepcopy(self)
            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

    def occupying(self, piece):
        return [sq for sq in squares if self.board[sq] == piece]

    def castle(self, side):
        color = self.turn
        rank = home_ranks[self.turn]

        if side not in ("queen", "king"):
            return False
        if not self.can_castle[color][side]:
            return False

        if side == "queen":
            check_if_empty    = cross("bcd", rank)
            check_if_attacked = cross("cde", rank)
            rook_move = "a"+rank, "d"+rank
            king_move = "e"+rank, "c"+rank
        elif side == "king":
            check_if_empty    = cross("fg", rank)
            check_if_attacked = check_if_empty
            rook_move = "h"+rank, "f"+rank
            king_move = "e"+rank, "g"+rank
        are_empty    = all(map(self.is_empty, check_if_empty))
        not_attacked = not any(self.is_attacked(color, sq)
                               for sq in check_if_attacked)

        if not (are_empty and not_attacked):
            return False

        self.board[king_move[1]] = self.board[king_move[0]]
        self.board[rook_move[1]] = self.board[rook_move[0]]
        self.board[king_move[0]] = Piece()
        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):
        new = deepcopy(self)
        if not new.is_legal(source, target, promotion):
            return False

        board = new.board
        moved = board[source]
        eaten = board[target]

        if promotion == None:
            promotion = moved
        else:
            promotion = Piece(new.turn, promotion)

        if moved.piece == "pawn" and target in moves(source, bishop_dirs) and eaten==None:
            # En passant
            a, b = move(target, up), move(target, down)
            board[a] = Piece()
            board[b] = Piece()

        # Castling
        if moved.piece == "rook":
            new.can_castle[new.turn][moved.side] = False
        if eaten.piece == "rook":
            new.can_castle[eaten.color][eaten.side] = False
        if moved.piece == "king":
            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

    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 in ("king", "pawn"):
            return False
        return color == self.turn and target in self.possible_moves(source)

    def is_attacked(self, color, sq):
        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 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
        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, back = up, down
            if color == "black":
                dir, back = down, up
            frwd = move(sq, dir)
            jump = move(frwd, dir)

            if frwd in squares and is_empty(frwd):
                r.append(frwd)
                is_on_pawn_rank = pawn_ranks[color] == sq[1]
                if is_on_pawn_rank and is_empty(jump):
                    r.append(jump)
            for t in targets:
                a, b = move(t, dir), move(t, back)
                en_passant = can_eat(b) and self.moves[-1] == (a,b)
                if can_eat(t) or en_passant:
                    r.append(t)
            return r
        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 []

        board = self.board
        piece = board[sq].piece
        color = board[sq].color

        def possible_line(dir):
            target = move(sq, dir)
            while target in board and self.is_empty(target):
                yield target
                target = move(target, dir)
            if target in board:
                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":
            if color == "white":
                dirs = upper_left, upper_right
            else:
                dirs = lower_left, lower_right
            return moves(sq, dirs)
        elif piece == "knight":
            nums = [2, -2, 1, -1]
            dirs = [(a, b) for a,b in product(nums, nums)
                               if abs(a) != abs(b)]
            return moves(sq, dirs)
        elif piece == "rook":
            return rook
        elif piece == "bishop":
            return bishop
        elif piece == "queen":
            return rook + bishop
        elif piece == "king":
            dirs = rook_dirs + bishop_dirs
            return moves(sq, dirs)
        return []



def test():
    game = 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()