140 lines
4.1 KiB
Python
140 lines
4.1 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)
|
|
|
|
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 addv(sq, v):
|
|
return chr(ord(sq[0]) + v[0]) + str(int(sq[1]) + v[1])
|
|
|
|
def addvs(sq, vectors):
|
|
return [addv(sq, v) for v in vectors]
|
|
|
|
def negatevs(vectors):
|
|
return [(-a, -b) for a,b in vectors]
|
|
|
|
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()
|
|
|
|
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 possible_moves(self, sq):
|
|
def is_possible(target):
|
|
return target in squares and \
|
|
self.board[target].color != self.board[sq].color
|
|
def are_possible(vecs):
|
|
r = []
|
|
sqrs = addvs(sq, vecs)
|
|
for s in sqrs:
|
|
if not is_possible(s):
|
|
break
|
|
r.append(s)
|
|
if not self.is_empty(s):
|
|
break
|
|
return r
|
|
|
|
def possible_rook():
|
|
rvecs = list(zip(range(1,8), [0]*7))
|
|
fvecs = [(b,a) for a,b in rvecs]
|
|
vecs = [rvecs, fvecs, negatevs(rvecs), negatevs(fvecs)]
|
|
return sum(map(are_possible, vecs), [])
|
|
|
|
def possible_bishop():
|
|
ldvecs = list(zip(range(1,8), range(1,8)))
|
|
rdvecs = list(zip(range(1,8), range(-1, -8, -1)))
|
|
vecs = [ldvecs, rdvecs, negatevs(ldvecs), negatevs(rdvecs)]
|
|
return sum(map(are_possible, vecs), [])
|
|
|
|
piece = self.board[sq].piece
|
|
|
|
if piece == "pawn":
|
|
pass
|
|
elif piece == "knight":
|
|
vecs = list(product((1,-1), (2,-2)))
|
|
vecs += list(product((2,-2), (1,-1)))
|
|
return [sq for sq in addvs(sq, vecs) if is_possible(sq)]
|
|
elif piece == "rook":
|
|
return possible_rook()
|
|
elif piece == "bishop":
|
|
return possible_bishop()
|
|
elif piece == "queen":
|
|
return possible_bishop() + possible_rook()
|
|
|
|
|
|
|
|
def test():
|
|
game = Game()
|
|
assert len(squares) == 8**2
|
|
assert sum(map(len, init_positions.values())) == 8*4
|
|
game.board["d4"] = Piece("black", "bishop")
|
|
print(game)
|
|
print (game.possible_moves("d4"))
|
|
|
|
test()
|