432 lines
13 KiB
Python
432 lines
13 KiB
Python
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, None)
|
|
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 ""
|
|
name = self.color.upper() + " CHESS " + self.piece.upper()
|
|
return lookup(name)
|
|
|
|
def __str__(self):
|
|
return repr(self)
|
|
|
|
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 endofgame(self):
|
|
return self.is_stall() or self.is_mate()
|
|
|
|
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)
|
|
self.turn = invert(self.turn)
|
|
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 validate_move(self, source, dest, promotion):
|
|
new = deepcopy(self)
|
|
valid = new.move(source, dest, promotion)
|
|
new.turn = self.turn
|
|
return valid, new
|
|
|
|
def AN_move(self, move):
|
|
info = parse_AN(move)
|
|
if info == False:
|
|
return False
|
|
if info.queenside:
|
|
return self.castle("queen")
|
|
if info.kingside:
|
|
return self.castle("king")
|
|
|
|
piece = info.piece
|
|
sqrs = self.occupying(Piece(self.turn, piece))
|
|
promotion = info.promotion
|
|
dest = info.dest
|
|
if info.file != None:
|
|
sqrs = [sq for sq in sqrs if sq[0]==info.file]
|
|
if info.rank != None:
|
|
sqrs = [sq for sq in sqrs if sq[1]==info.rank]
|
|
moves = []
|
|
for sq in sqrs:
|
|
valid, new = self.validate_move(sq, dest, promotion)
|
|
check = new.is_check()==info.check and new.is_mate()==info.mate
|
|
if valid and check:
|
|
moves.append((sq, dest, promotion))
|
|
if len(moves) != 1:
|
|
return False
|
|
move = moves[0]
|
|
self.move(*move)
|
|
return True
|
|
|
|
|
|
|
|
def test():
|
|
game = Game()
|
|
while not game.is_stall():
|
|
move = input(">> ")
|
|
v = game.AN_move(move)
|
|
if not v:
|
|
print("Invalid")
|
|
continue
|
|
print(game)
|
|
|