Further refactoring of possible_moves and helper functions
parent
0b16fabcd8
commit
dad2fc7d58
125
chess.py
125
chess.py
|
@ -21,17 +21,32 @@ init_positions = {"pawn" : cross(files, pawn_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)
|
||||
|
||||
up_left, up_right = (-1, 1), (1, 1)
|
||||
down_left, down_right = (-1, -1), (1, -1)
|
||||
bishop_dirs = (up_left, up_right,
|
||||
down_left, down_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):
|
||||
|
@ -109,10 +124,50 @@ class Game:
|
|||
def is_legal(self, source, target):
|
||||
return target in self.possible_moves(source)
|
||||
|
||||
def is_attacked(self, sq):
|
||||
return False
|
||||
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 []
|
||||
|
||||
|
@ -120,21 +175,12 @@ class Game:
|
|||
piece = board[sq].piece
|
||||
color = board[sq].color
|
||||
|
||||
def can_eat(target):
|
||||
if target not in board:
|
||||
return False
|
||||
eaten = board[target]
|
||||
return eaten != None and color != eaten.color
|
||||
|
||||
def can_move_to(sq):
|
||||
return sq in board and self.is_empty(sq)
|
||||
|
||||
def possible_line(dir):
|
||||
target = move(sq, dir)
|
||||
while can_move_to(target):
|
||||
while target in board and self.is_empty(target):
|
||||
yield target
|
||||
target = move(target, dir)
|
||||
if can_eat(target):
|
||||
if target in board:
|
||||
yield target
|
||||
|
||||
def possible_lines(dirs):
|
||||
|
@ -144,31 +190,16 @@ class Game:
|
|||
bishop = possible_lines(bishop_dirs)
|
||||
|
||||
if piece == "pawn":
|
||||
r = []
|
||||
dir = up
|
||||
if color == "black":
|
||||
dir = down
|
||||
frwd = move(sq, dir)
|
||||
jump = move(frwd, dir)
|
||||
eat = (move(frwd, left),
|
||||
move(frwd, right))
|
||||
|
||||
if can_move_to(frwd):
|
||||
r.append(frwd)
|
||||
is_on_pawn_rank = (color == "black" and sq[1]=='7') \
|
||||
or \
|
||||
(color == "white" and sq[1]=='2')
|
||||
if can_move_to(jump) and is_on_pawn_rank:
|
||||
r.append(jump)
|
||||
return r + [sq for sq in eat if can_eat(sq)]
|
||||
|
||||
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]
|
||||
vectors = [(a, b) for a,b in product(nums, nums)
|
||||
if abs(a) != abs(b)]
|
||||
targets = [move(sq, v) for v in vectors]
|
||||
return [sq for sq in targets
|
||||
if can_move_to(sq) or can_eat(sq)]
|
||||
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":
|
||||
|
@ -176,10 +207,8 @@ class Game:
|
|||
elif piece == "queen":
|
||||
return rook + bishop
|
||||
elif piece == "king":
|
||||
vectors = rook_dirs + bishop_dirs
|
||||
targets = [move[sq, v] for v in vectors]
|
||||
return [sq for sq in targets
|
||||
if can_move_to(sq) and not is_attacked(sq)]
|
||||
dirs = rook_dirs + bishop_dirs
|
||||
return moves(sq, dirs)
|
||||
return []
|
||||
|
||||
|
||||
|
@ -188,9 +217,15 @@ def test():
|
|||
game = Game()
|
||||
assert len(squares) == 8**2
|
||||
assert sum(map(len, init_positions.values())) == 8*4
|
||||
moves = [("a2", "a3"), ("b1", "c3"), ("d1", "d2")]
|
||||
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()
|
||||
|
|
Loading…
Reference in New Issue