diff --git a/chess.py b/chess.py index 2ebcdf3..639b0a2 100644 --- a/chess.py +++ b/chess.py @@ -3,7 +3,7 @@ from unicodedata import lookup def cross(A, B): return tuple(a+b for a in A for b in B) -ranks = "12345678" +ranks = "87654321" files = "abcdefgh" squares = cross(files, ranks) @@ -29,6 +29,13 @@ class Piece: 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() @@ -63,7 +70,10 @@ class Game: self.board[sq].color = "black" def is_empty(self, sq): - return self.board[sq].piece == None + return self.board[sq] == None + + def occupying(self, piece): + return [sq for sq in squares if self.board[sq] == piece] @@ -73,5 +83,6 @@ def test(): assert len(squares) == 8**2 assert sum(map(len, init_positions.values())) == 8*4 print(game) + print(game.occupying(Piece("white", "pawn"))) test()