Allow providing board as context to Entry.move() (#398)

chess/__init__.py

@@ -3319,7 +3319,7 @@ def generate_castling_moves(self, from_mask: Bitboard = BB_ALL, to_mask: Bitboar
                 yield self._from_chess960(self.chess960, msb(king), candidate)
 
     def _from_chess960(self, chess960: bool, from_square: Square, to_square: Square, promotion: Optional[PieceType] = None, drop: Optional[PieceType] = None) -> Move:
-        if not chess960 and drop is None:
+        if not chess960 and promotion is None and drop is None:
             if from_square == E1 and self.kings & BB_E1:
                 if to_square == H1:
                     return Move(E1, G1)

chess/polyglot.py

@@ -305,7 +305,7 @@ class Entry(collections.namedtuple("Entry", "key raw_move weight learn")):
 
     __slots__ = ()
 
-    def move(self, *, chess960: bool = False) -> chess.Move:
+    def move(self, *, board: chess.Board = None) -> chess.Move:
         """Gets the move (as a :class:`~chess.Move` object)."""
         # Extract source and target square.
         to_square = self.raw_move & 0x3f
@@ -315,8 +315,18 @@ def move(self, *, chess960: bool = False) -> chess.Move:
         promotion_part = (self.raw_move >> 12) & 0x7
         promotion = promotion_part + 1 if promotion_part else None
 
-        # Convert castling moves.
-        if not chess960 and not promotion:
+        # Piece drop.
+        if from_square == to_square:
+            drop, promotion = promotion, None
+        else:
+            drop = None
+
+        # Normalize castling moves if we have a board as context.
+        if board is not None:
+            return board._from_chess960(board.chess960, from_square, to_square, promotion, drop)
+
+        # Best effort to normalize castling moves without context.
+        if not promotion and not drop:
             if from_square == chess.E1:
                 if to_square == chess.H1:
                     return chess.Move(chess.E1, chess.G1)
@@ -328,10 +338,7 @@ def move(self, *, chess960: bool = False) -> chess.Move:
                 elif to_square == chess.A8:
                     return chess.Move(chess.E8, chess.C8)
 
-        if promotion and from_square == to_square:
-            return chess.Move(from_square, to_square, drop=promotion)
-        else:
-            return chess.Move(from_square, to_square, promotion)
+        return chess.Move(from_square, to_square, promotion, drop)
 
 
 class MemoryMappedReader:
@@ -419,7 +426,7 @@ def find_all(self, board: Union[chess.Board, int], *, minimum_weight: int = 1, e
                 continue
 
             if context:
-                move = entry.move(chess960=context.chess960)
+                move = entry.move(board=context)
             elif exclude_moves:
                 move = entry.move()
 

test.py

@@ -1789,6 +1789,12 @@ def test_castling(self):
             self.assertIn(pos.parse_san("O-O-O"), moves)
             self.assertEqual(len(moves), 1)
 
+            # Not a castling move.
+            pos = chess.Board("1r1qr1k1/1b2bp1n/p2p2pB/1pnPp2p/P1p1P3/R1P2NNP/1PBQ1PP1/4R1K1 w - - 0 1")
+            entry = book.find(pos)
+            self.assertEqual(entry.move(), chess.Move.from_uci("e1c1"))
+            self.assertEqual(entry.move(board=pos), chess.Move.from_uci("e1a1"))
+
     def test_empty_book(self):
         with chess.polyglot.open_reader(os.devnull) as book:
             self.assertEqual(len(book), 0)