Single_Player_Chess/chess.py

from itertools import product
from unicodedata import lookup

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 = "27"
home_ranks = "18"
init_positions = {"pawn"   : cross(files, pawn_ranks),
                  "rook"   : cross("ah",  home_ranks),
                  "knight" : cross("bg",  home_ranks),
                  "bishop" : cross("cf",  home_ranks),
                  "queen"  : cross("d",   home_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)

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 = []

    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"

    def is_empty(self, sq):
        return self.board[sq] == None

    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):
            return False
        board = self.board
        moved = board[source]
        eaten = board[target]

        self.moves.append((source, target))
        self.stack.append(eaten)

        board[source] = Piece()
        board[target] = moved
        return True

    def is_legal(self, source, target):
        return 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 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 []

        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()
    assert len(squares) == 8**2
    assert sum(map(len, init_positions.values())) == 8*4
    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()