Added doc strings to chess.py

master
Tibor Bizjak 2019-09-10 23:27:31 +02:00
parent 661e87715a
commit f402ae36af
1 changed files with 68 additions and 21 deletions

View File

@ -3,6 +3,7 @@ from unicodedata import lookup
from copy import deepcopy
def cross(A, B):
"""Return sum of of every element from A and B."""
return tuple(a+b for a in A for b in B)
pieces = ["pawn", "knight", "bishop", "rook", "queen", "king"]
@ -33,22 +34,29 @@ AN_names = {'R' : "rook",
'K' : "king"}
def move(sq, v):
"""Add vector v to square and return a square, even if not in board."""
return chr(ord(sq[0]) + v[0]) + str(int(sq[1]) + v[1])
def moves(sq, dirs):
"""Add multiple vectors to a square. Return squares that are in board."""
targets = [move(sq, dir) for dir in dirs]
return [sq for sq in targets if sq in squares]
def invert(color):
"""Return inverted color."""
if color == "white":
return "black"
return "white"
class Empty:
"""Empty class to avoid unnecessary quotations."""
def __repr__(self):
return str(vars(self))
def parse_AN(move):
"""Parse chess move writen in algebraic notation. Return relevant
information encapsulated within an Empty class, or False if parsing
fails."""
def suffix_in(l, dval=False, i=1):
nonlocal move
x = move[-i:]
@ -87,6 +95,7 @@ def parse_AN(move):
r.piece = AN_names[r.piece]
return r
# Names for vectors
up, down = (0, 1), (0, -1)
left, right = (-1, 0), (1, 0)
rook_dirs = (up, down, left, right)
@ -97,20 +106,27 @@ bishop_dirs = (upper_left, upper_right,
lower_left, lower_right)
class Piece:
"""Contains piece color and piece type"""
def __init__(self, color=None, piece=None):
"""Assign color and piece. If not specified they
default to None."""
self.color = color
self.piece = piece
def __repr__(self):
"""Returns unicode representation of instance."""
if None in (self.piece, self.color):
return ""
name = self.color.upper() + " CHESS " + self.piece.upper()
return lookup(name)
def __str__(self):
"""Same as repr."""
return repr(self)
def __eq__(self, other):
"""Ensures equality testing. If color or piece or both
equal None, instance equals None."""
if other == None:
return None in (self.color, self.piece)
if not isinstance(other, Piece):
@ -118,6 +134,16 @@ class Piece:
return self.piece == other.piece and self.color == other.color
class Game:
"""Chess game class.
Squares are represented by strings made of letters and integers e.g. 'h4'.
Board is represented by a dictionary whose keys are squares
and values instances of Piece class.
Moves are represented internally by tuples of squares;
a source and destination.
Moves are appended to moves stack.
Upon each move, the relevant attributes of the instance (the game state)
are appended to the history stack. This allows loading previous
game states to instance variables."""
def __init__(self):
self.index = 0
self.make_board()
@ -135,7 +161,7 @@ class Game:
self.turn))
def __repr__(self):
# Unicode board representation
"""Return unicode board representation."""
r = ""
for rank in ranks:
r += rank + " |"
@ -152,6 +178,7 @@ class Game:
return r
def make_board(self):
"""Create board and store it in self.board."""
self.board = dict((sq, Piece()) for sq in squares)
# Add pieces
for piece, positions in init_positions.items():
@ -169,10 +196,15 @@ class Game:
self.board["h8"].side = "king"
def rebase(self):
"""Shorten moves and history stacks if move is made
while past game state is loaded into memory. Essentialy
overrides future game states."""
self.moves = self.moves[:self.index+1]
self.history = self.history[:self.index+1]
def timetravel(self, i):
"""Load game state from past into instance variables.
Can be reversed since history and moves stack stay intact."""
if i >= len(self.history) or i < 0:
return False
c, board, turn = self.history[i]
@ -182,30 +214,39 @@ class Game:
self.index = i
def prev(self):
"""Load game state from previous turn, if it exists."""
return self.timetravel(self.index-1)
def next(self):
"""Load game state from next turn, if it exists."""
return self.timetravel(self.index+1)
def last(self):
"""Load game state from last move."""
return self.timetravel(len(self.history)-1)
def first(self):
"""Load first recorded game state, an empty board."""
return self.timetravel(0)
def is_empty(self, sq):
"""Return True if square is empty, else False."""
return self.board[sq] == None
def is_check(self, color=None):
"""Return True if king of active player is in check else False."""
if color == None:
color = self.turn
sq = self.occupying(Piece(color, "king"))[0]
return self.is_attacked(color, sq)
def is_mate(self):
"""Return True if checkmate for active player."""
return self.is_check() and self.is_stall()
def is_stall(self):
"""Return True if active player has no legal move. Whether
the king is in check is irrelevant."""
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)]:
@ -215,12 +256,16 @@ class Game:
return True
def endofgame(self):
"""Return True if active player can't make a legal move."""
return self.is_stall() or self.is_mate()
def occupying(self, piece):
"""Return squares occupied by piece. Piece variable is an
instance of Piece class."""
return [sq for sq in squares if self.board[sq] == piece]
def castle(self, side):
"""Castle if possible, if not return False."""
color = self.turn
rank = home_ranks[self.turn]
@ -262,6 +307,9 @@ class Game:
return True
def move(self, source, target, promotion=None):
"""Move piece from source to target if legal and promote it
to promotion if it's a pawn whos reached the last
rank. Return False if fails."""
new = deepcopy(self)
if not new.is_legal(source, target, promotion):
return False
@ -308,6 +356,9 @@ class Game:
return True
def is_legal(self, source, target, promotion):
"""Return True if move is partially legal. Does not
check whether the active players king is in check
after the move! See the validate_move method."""
piece = self.board[source].piece
color = self.board[source].color
pawn_on_last_rank = piece == "pawn" and \
@ -321,6 +372,7 @@ class Game:
return color == self.turn and target in self.possible_moves(source)
def is_attacked(self, color, sq):
"""Return True if square is attacked by enemy of opposite color."""
enemy_color = invert(color)
attacked = []
for piece in pieces:
@ -329,6 +381,9 @@ class Game:
return sq in attacked
def all_legal_moves(self):
"""Return all possible partially legal moves by active
player. These moves can put the active player in check,
thus they are not completely legal. Castlings not included."""
color = self.turn
sqrs = []
for piece in pieces:
@ -341,6 +396,9 @@ class Game:
return r
def possible_moves(self, sq):
"""Return all squares the piece from sq can move to.
Moving to these squares might put the player in check.
Castlings not included."""
board = self.board
piece = board[sq].piece
color = board[sq].color
@ -380,6 +438,8 @@ class Game:
return targets
def attacks(self, sq):
"""Return all squares the piece on sq attacks. Return empty list
if square is empty."""
if sq not in squares:
return []
@ -424,12 +484,19 @@ class Game:
return []
def validate_move(self, source, dest, promotion):
"""Create a copy of the game and try given move.
Return True if the move succeeds. This method
exists because is_legal does not test whether
the move puts the player in check."""
new = deepcopy(self)
valid = new.move(source, dest, promotion)
new.turn = self.turn
return valid, new
def AN_move(self, move):
"""Move is given in algebraic notation. The method
moves the piece if the move is legal and unambiguous,
else it returns False."""
AN = move
info = parse_AN(move)
if info == False:
@ -462,23 +529,3 @@ class Game:
def test():
game = Game()
move = ""
while not game.is_stall() and move!='q':
move = input(">> ")
if move == "prev":
game.prev()
print(game)
continue
elif move == "next":
game.next()
print(game)
continue
v = game.AN_move(move)
if not v:
print("Invalid")
continue
print(game)
print (game.history)