Implement castling and create a utility function for applying move on a board and taking into account promotion during legal move generation

app.py

@@ -164,6 +164,32 @@ class Board:
     def __str__(self):
         return "\n".join([" ".join([p.value for p in self.grid[i]]) for i in range(8)])
 
+    def make_move(self, src: str, dst: str):
+        sx, sy = pos_to_coord(src)
+        dx, dy = pos_to_coord(dst)
+        self.grid[dy][dx] = self.index(src)
+        piece_kind = self.grid[sy][sx].value.lower()
+        is_white = self.grid[sy][sx].value.isupper()
+
+        if piece_kind == "p" and (dy == 7 or dy == 0):
+            self.grid[dy][dx] = Piece.WHITE_QUEEN if is_white else Piece.BLACK_QUEEN
+
+        if (
+            is_white and dst == "e1" and piece_kind == "k" and dst == "g1"
+        ):  # Castling (White)
+            castle_from = pos_to_coord("h1")
+            castle_to = pos_to_coord("f1")
+            self.grid[castle_from.y][castle_from.x] = Piece.EMPTY
+            self.grid[castle_to.y][castle_to.x] = Piece.WHITE_CASTLE
+        if (
+            not is_white and src == "e8" and piece_kind == "k" and dst == "g8"
+        ):  # Castling (Black)
+            castle_from = pos_to_coord("h8")
+            castle_to = pos_to_coord("f8")
+            self.grid[castle_from.y][castle_from.x] = Piece.EMPTY
+            self.grid[castle_to.y][castle_to.x] = Piece.BLACK_CASTLE
+        self.grid[sy][sx] = Piece.EMPTY
+
 
 class Room:
     board: Board
@@ -434,6 +460,20 @@ def get_piece_moves(piece_kind, board: Board, is_white, src: str) -> list[Coord]
     elif piece_kind == "k":
         x, y = pos_to_coord(src)
         comb = product([-1, 0, 1], repeat=2)
+        if (
+            src == "e1"
+            and is_white
+            and board.index("f1") == Piece.EMPTY
+            and board.index("g1") == Piece.EMPTY
+        ):  # Castle (White)
+            valids.append(pos_to_coord("g1"))
+        if (
+            src == "e8"
+            and is_white
+            and board.index("f8") == Piece.EMPTY
+            and board.index("g8") == Piece.EMPTY
+        ):  # Castling (Black)
+            valids.append(pos_to_coord("g8"))
         for dir in comb:
             if dir[0] ** 2 + dir[1] ** 2 == 0:
                 continue  # x=0 y=0 cannot move => invalid
@@ -476,9 +516,7 @@ def generate_valid_moves(
                 break
     for m in possible_moves:
         fake_board = board.copy()
-        fake_board.grid[m.y][m.x] = fake_board.index(src)
-        sx, sy = pos_to_coord(src)
-        fake_board.grid[sy][sx] = Piece.EMPTY
+        fake_board.make_move(src, coord_to_pos_safe(m))
         king_safe = True
         for i in range(8):
             for j in range(8):
@@ -551,12 +589,8 @@ async def move(request: Request, room_id):
     ):
         return 400, {"code": "HTME", "error": "Cannot move to your own piece"}
     valid_moves = generate_valid_moves(piece_kind, board, is_white, src)
-    if (c := pos_to_coord(dst)) in valid_moves:
-        board.grid[c.y][c.x] = srcp
-        sx, sy = pos_to_coord(src)
-        board.grid[sy][sx] = Piece.EMPTY
-        if (c.y == 0 or c.y == 7) and piece_kind == "p":
-            board.grid[c.y][c.x] = Piece.WHITE_QUEEN if is_white else Piece.BLACK_QUEEN
+    if pos_to_coord(dst) in valid_moves:
+        board.make_move(src, dst)
         room.turn = Color.BLACK if color == Color.WHITE else Color.WHITE
         room.last_move = datetime.now(timezone.utc)
         opp_checkmate = True