Initial commit

cleaning/format.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+"""
+Turn refined PGN games (Site/WhiteElo/BlackElo headers + movetext with
+%eval comments) into a CSV where each row is the board state *after* one
+move.
+
+Row layout (67 columns):
+    game_id, a8, b8, c8, d8, e8, f8, g8, h8, a7, ... , h1, eval, turn
+
+  - game_id : 0-based index of the game in the input file. Not asked for
+              explicitly, but with many games concatenated into one CSV
+              there is no other way to tell which rows belong to which
+              game, so I added it as column 1. Drop it with --no-game-id
+              if you don't want it.
+  - a8..h1  : the 64 squares (rank 8 down to rank 1, file a to h), each
+              holding the FEN piece letter on that square after the move
+              (uppercase = white, lowercase = black) or an empty string
+              if the square is empty.
+  - eval    : the %eval attached to that move (blank if that particular
+              move had none).
+  - turn    : "white" or "black" -- whose turn it is to move NEXT.
+
+Streams game-by-game (like the two earlier scripts) so it stays cheap on
+RAM no matter how large the input file is. Requires `pip install chess`.
+
+Usage:
+    python pgn_to_csv.py refined.pgn out.csv
+    python pgn_to_csv.py refined.pgn out.csv --no-game-id
+"""
+
+import argparse
+import csv
+import re
+import sys
+
+import chess
+
+TAG_RE = re.compile(r'^\[(\w+)\s+"(.*)"\]\s*$')
+EVAL_RE = re.compile(r'\[%eval\s+([^\]]+)\]')
+MOVE_NUM_RE = re.compile(r'^\d+\.+$')
+RESULT_RE = re.compile(r'^(1-0|0-1|1/2-1/2|\*)$')
+# Strip trailing annotation glyphs python-chess's SAN parser won't accept
+# (e.g. "e5?!", "Nc3!!", "Qxd1+!") while keeping the "+"/"#" that ARE part
+# of legal SAN.
+SUFFIX_ANNOTATION_RE = re.compile(r'[!?]+$')
+
+SQUARE_COLUMNS = [f + str(r) for r in range(8, 0, -1) for f in "abcdefgh"]
+
+
+def tokenize_movetext(movetext: str):
+    """Yield (san_move, eval_str_or_None) pairs from a movetext blob."""
+    tokens = re.findall(r'\{[^}]*\}|\S+', movetext)
+    i = 0
+    n = len(tokens)
+    while i < n:
+        tok = tokens[i]
+        if MOVE_NUM_RE.match(tok) or RESULT_RE.match(tok):
+            i += 1
+            continue
+        move_san = tok
+        eval_val = None
+        if i + 1 < n and tokens[i + 1].startswith('{'):
+            m = EVAL_RE.search(tokens[i + 1])
+            if m:
+                eval_val = m.group(1)
+            i += 2
+        else:
+            i += 1
+        yield move_san, eval_val
+
+
+def board_row(board: chess.Board):
+    """Return the 64 square values for the current board, a8..h1 order."""
+    row = []
+    for name in SQUARE_COLUMNS:
+        sq = chess.parse_square(name)
+        piece = board.piece_at(sq)
+        row.append(piece.symbol() if piece else "")
+    return row
+
+
+def process_game(game_id, header_lines, movetext_lines, writer, include_game_id):
+    if not movetext_lines:
+        return 0
+
+    movetext = ' '.join(line.strip() for line in movetext_lines if line.strip())
+
+    board = chess.Board()
+    rows_written = 0
+
+    for move_san, eval_val in tokenize_movetext(movetext):
+        clean_san = SUFFIX_ANNOTATION_RE.sub('', move_san)
+        try:
+            board.push_san(clean_san)
+        except (chess.IllegalMoveError, chess.InvalidMoveError, chess.AmbiguousMoveError, ValueError) as e:
+            print(f"  [game {game_id}] stopping early: could not parse move "
+                  f"'{move_san}' ({e})", file=sys.stderr)
+            break
+
+        turn = "white" if board.turn == chess.WHITE else "black"
+        row = (board_row(board) + [eval_val if eval_val is not None else "", turn])
+        if include_game_id:
+            row = [game_id] + row
+        writer.writerow(row)
+        rows_written += 1
+
+    return rows_written
+
+
+def refine(input_path: str, output_path: str, include_game_id: bool):
+    header_lines = []
+    movetext_lines = []
+    state = "seek_header"
+    game_id = 0
+    total_rows = 0
+
+    header_row = SQUARE_COLUMNS + ["eval", "turn"]
+    if include_game_id:
+        header_row = ["game_id"] + header_row
+
+    with open(input_path, "r", encoding="utf-8", errors="replace", buffering=1024 * 1024) as fin, \
+         open(output_path, "w", encoding="utf-8", newline="") as fout:
+
+        writer = csv.writer(fout)
+        writer.writerow(header_row)
+
+        for raw_line in fin:
+            line = raw_line.rstrip("\n")
+
+            if state == "seek_header":
+                if line.startswith("["):
+                    header_lines.append(line)
+                elif line.strip() == "" and header_lines:
+                    state = "seek_movetext"
+            elif state == "seek_movetext":
+                if line.strip() == "":
+                    total_rows += process_game(game_id, header_lines, movetext_lines,
+                                                writer, include_game_id)
+                    game_id += 1
+                    header_lines = []
+                    movetext_lines = []
+                    state = "seek_header"
+                else:
+                    movetext_lines.append(line)
+
+        if header_lines and movetext_lines:
+            total_rows += process_game(game_id, header_lines, movetext_lines,
+                                        writer, include_game_id)
+            game_id += 1
+
+    return game_id, total_rows
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert refined PGN games into a per-move board-state CSV.")
+    parser.add_argument("input", help="Path to the refined PGN file")
+    parser.add_argument("output", help="Path to write the CSV file")
+    parser.add_argument("--no-game-id", action="store_true",
+                         help="Omit the leading game_id column (66-col strict layout).")
+    args = parser.parse_args()
+
+    games, rows = refine(args.input, args.output, include_game_id=not args.no_game_id)
+    print(f"Processed {games} games -> {rows} rows written to {args.output}")
+
+
+if __name__ == "__main__":
+    main()