232 lines
6.6 KiB
Python
232 lines
6.6 KiB
Python
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()
|