PGN4 to PGN conversion

pgn4_to_pgn.py

import argparse
import os
import re
import sys

import pyffish


SEIRAWAN_CASTLING = {'O-O': 'O-O', 'O-O-O': 'O-O-O', 'h1-e1': 'O-O', 'h8-e8': 'O-O', 'a1-e1': 'O-O-O', 'a8-e8': 'O-O-O'}


def coords(match, files, ranks):
    return chr(ord(match.group(1)) - (11 - files)) + str(int(match.group(2)) - (11 - ranks))


def gating(match):
    return (SEIRAWAN_CASTLING[match.group(1)] + '/' + match.group(2) + match.group(3) if match.group(1) in SEIRAWAN_CASTLING
            else match.group(1) + '/' + match.group(2))


def pgn4_to_pgn(pgn4, override_variant, files, ranks):
    pgn4 = re.sub(r' .. ', ' ', pgn4)
    pgn4 = re.sub(r'x[A-Z]([a-l][0-9]{1,2})', r'x\1', pgn4)  # captures
    pgn4 = re.sub(r'([a-l])([0-9]{1,2})', lambda x: coords(x, files, ranks), pgn4)  # coords
    pgn4 = re.sub(r'([^ ]*)&@[a-z]([A-Z])-([a-l][0-9]{1,2})', gating, pgn4)  # gating

    pgn = ''
    variant = override_variant
    start_fen = pyffish.start_fen(override_variant) if override_variant else None
    for line in pgn4.splitlines():
        if line.startswith('['):
            moves = []
            if line.startswith('[Variant'):
                continue
            if line.startswith('[Site'):
                if not override_variant:
                    variant = re.search(r'variants\/([^/]*)', line).group(1).replace("-chess", "").replace('-', '')
                    if variant not in pyffish.variants():
                        raise Exception('Unsupported variant {}'.format(variant))
                    start_fen = pyffish.start_fen(variant)
                line = '[Variant "{}"]'.format(variant.capitalize())
            pgn += line + os.linesep
        else:
            parts = line.split()
            for i, move in enumerate(parts):
                fen = pyffish.get_fen(variant, start_fen, moves)
                legals = pyffish.legal_moves(variant, fen, [])
                lan_legals = {pyffish.get_san(variant, fen, m, False, pyffish.NOTATION_LAN): m for m in legals}
                if move[0].isdigit() or move in ('T', 'R', 'S'):
                    continue
                else:
                    # fix weird game end markers
                    move = move.rstrip('S').replace('+#', '#')
                    if move in lan_legals:
                        moves.append(lan_legals[move])
                        parts[i] = pyffish.get_san(variant, fen, lan_legals[move], False, pyffish.NOTATION_SAN)
                    else:
                        sys.stderr.write('Warning: Move {} not found\n'.format(move))
            pgn += ' '.join(parts) + os.linesep
    return pgn


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('pgn4', nargs='+')
    parser.add_argument('--ranks', type=int, default=8)
    parser.add_argument('--files', type=int, default=8)
    parser.add_argument('--variant', choices=pyffish.variants())
    args = parser.parse_args()

    for pgn in args.pgn4:
        with open(pgn) as f:
            pgn4 = f.read()
            pgn = pgn4_to_pgn(pgn4, args.variant, args.files, args.ranks)
            print(pgn)