pgn.py

1"""Bounded, fail-closed PGN lexical validation and semantic parsing."""
2import contextlib
3import io
4import re
5import subprocess
6import chess
7from .tokenizer import State, BOS, EOS
8
9RESULTS = {"1-0", "0-1", "1/2-1/2", "*"}
10HEADER = re.compile(r'^\[([A-Za-z][A-Za-z0-9_]*)\s+"((?:[^"\\]|\\["\\])*)"\s*\]$')
11SAN = re.compile(r'^(?:O-O(?:-O)?|0-0(?:-0)?|[KQRBN]?[a-h]?[1-8]?x?[a-h][1-8](?:=[QRBN])?)[+#]?[!?]{0,2}$')
12
13
14@contextlib.contextmanager
15def open_pgn(path):
16    """Ordinary zstd frames stream from byte zero; decoder exit is checked."""
17    if str(path).endswith('.zst'):
18        with open(path, 'rb') as source, contextlib.ExitStack() as stack:
19            import tempfile
20            err = stack.enter_context(tempfile.TemporaryFile())
21            proc = subprocess.Popen(['zstd', '-dc'], stdin=source, stdout=subprocess.PIPE, stderr=err)
22            stream = io.TextIOWrapper(proc.stdout, encoding='utf-8-sig', errors='strict', newline=None)
23            try:
24                yield stream
25                stream.close()
26                if proc.wait() != 0:
27                    raise ValueError('zstd decoder failed (corrupt/truncated source)')
28            finally:
29                stream.close()
30                if proc.poll() is None:
31                    proc.terminate()
32                proc.wait()
33    else:
34        with open(path, encoding='utf-8-sig', errors='strict') as stream:
35            yield stream
36
37
38def lex(stream, max_atom=8192):
39    """Discard comments without buffering; all other atoms are size bounded."""
40    def chars():
41        while chunk := stream.read(65536):
42            yield from chunk
43    iterator = iter(chars())
44    char = next(iterator, '')
45    line_start = True
46    while char:
47        if char.isspace():
48            line_start = char == '\n' or (line_start and char in ' \t\r')
49            char = next(iterator, '')
50            continue
51        if char == ';' or (char == '%' and line_start):
52            while char and char != '\n':
53                char = next(iterator, '')
54            continue
55        line_start = False
56        if char == '{':
57            char = next(iterator, '')
58            while char and char != '}':
59                char = next(iterator, '')
60            if not char:
61                yield ('ERROR', 'unterminated_comment')
62            char = next(iterator, '')
63            continue
64        if char == '[':
65            atom, quote, escape, oversized = '[', False, False, False
66            char = next(iterator, '')
67            while char:
68                if len(atom) < max_atom:
69                    atom += char
70                else:
71                    oversized = True
72                if char == ']' and not quote:
73                    break
74                if char == '"' and not escape:
75                    quote = not quote
76                escape = char == '\\' and not escape
77                char = next(iterator, '')
78            yield ('HEADER' if char and not oversized else 'ERROR', atom if not oversized else 'oversized_header')
79            char = next(iterator, '')
80            continue
81        if char in '()':
82            yield ('ATOM', char)
83            char = next(iterator, '')
84            continue
85        atom, oversized = '', False
86        while char and not char.isspace() and char not in '[{;()':
87            if len(atom) < max_atom:
88                atom += char
89            else:
90                oversized = True
91            char = next(iterator, '')
92        if oversized:
93            yield ('ERROR', 'oversized_atom')
94        else:
95            # Move numbers may be glued to SAN: 1.e4, 1...e5.
96            match = re.match(r'^(\d+\.(?:\.\.)?)(.*)$', atom)
97            if match:
98                yield ('ATOM', match[1])
99                atom = match[2]
100            if atom:
101                yield ('ATOM', atom)
102
103
104def records(stream, max_record_bytes=1024 * 1024):
105    """Headers begin games. Headerless movetext is supported as one record.
106
107    Multiple headerless games without headers are rejected as one ambiguous record.
108    Oversized records are discarded until the next unambiguous header boundary.
109    Comments and variations are not used to infer game boundaries.
110    """
111    record, size, body, depth, overflow = [], 0, False, 0, False
112    for kind, value in lex(stream):
113        if kind == 'HEADER' and body and depth == 0:
114            yield record if not overflow else [('ERROR', 'oversized_record')]
115            record, size, body, overflow = [], 0, False, False
116        size += len(value.encode('utf-8')) + 16
117        if size <= max_record_bytes:
118            record.append((kind, value))
119        else:
120            overflow = True
121        if kind != 'HEADER':
122            body = True
123        if value == '(':
124            depth += 1
125        elif value == ')':
126            depth = max(0, depth - 1)
127    if record or overflow:
128        yield record if not overflow else [('ERROR', 'oversized_record')]
129
130
131class Rejected(ValueError):
132    pass
133
134
135def parse(record):
136    headers, atoms = {}, []
137    for kind, value in record:
138        if kind == 'ERROR':
139            raise Rejected(value if len(value) < 60 else 'lexical_error')
140        if kind == 'HEADER':
141            match = HEADER.fullmatch(value)
142            if not match or atoms or match[1] in headers or '\n' in value or '\r' in value:
143                raise Rejected('invalid_header')
144            headers[match[1]] = re.sub(r'\\(["\\])', r'\1', match[2])
145        else:
146            atoms.append(value)
147    if headers.get('Variant', 'Standard') not in ('Standard', 'Chess'):
148        raise Rejected('unsupported_variant')
149    if (headers.get('SetUp', '0') not in ('0', '1')
150            or (headers.get('SetUp') == '1' and 'FEN' not in headers)
151            or ('FEN' in headers and headers['FEN'] != chess.STARTING_FEN)):
152        raise Rejected('nonstandard_start')
153    if headers.get('Result', '*') not in RESULTS:
154        raise Rejected('invalid_result')
155    board, stack, moves, ending = chess.Board(), [], [], None
156    for atom in atoms:
157        if ending is not None:
158            raise Rejected('trailing_movetext')
159        if atom == '(':
160            if not board.move_stack:
161                raise Rejected('invalid_variation')
162            stack.append(board.copy(stack=1))
163            board.pop()
164        elif atom == ')':
165            if not stack:
166                raise Rejected('unbalanced_variation')
167            board = stack.pop()
168        elif re.fullmatch(r'\d+\.(?:\.\.)?', atom):
169            digits = atom.split('.')[0]
170            if len(digits) > 10:
171                raise Rejected('incorrect_move_number')
172            number = int(digits)
173            if number != board.fullmove_number or (atom.endswith('...') != (board.turn == chess.BLACK)):
174                raise Rejected('incorrect_move_number')
175        elif re.fullmatch(r'\$\d+', atom) or atom in ('!', '?', '!!', '??', '!?', '?!'):
176            continue
177        elif atom in RESULTS:
178            if stack:
179                raise Rejected('result_inside_variation')
180            ending = atom
181        else:
182            if not SAN.fullmatch(atom):
183                raise Rejected('invalid_san_or_null')
184            try:
185                move = board.parse_san(atom.rstrip('!?'))
186            except ValueError as error:
187                raise Rejected('illegal_or_ambiguous_move') from error
188            if not move or move not in board.legal_moves:
189                raise Rejected('illegal_or_null_move')
190            supplied = atom.rstrip('!?').replace('0', 'O')
191            canonical = board.san(move)
192            if supplied.rstrip('+#') != canonical.rstrip('+#') or (supplied.endswith(('+', '#')) and supplied[-1] != canonical[-1]):
193                raise Rejected('noncanonical_or_false_check_san')
194            board.push(move)
195            if not stack:
196                moves.append(move)
197    if stack or ending is None:
198        raise Rejected('incomplete_game')
199    if 'Result' in headers and ending != headers['Result']:
200        raise Rejected('result_mismatch')
201    headers['Result'] = ending
202    state, tokens = State(), [BOS]
203    state.feed(BOS)
204    for move in moves:
205        tokens.extend(state.append_move(move))
206    tokens.append(EOS)
207    state.feed(EOS)
208    restored = State()
209    recovered = []
210    for token in tokens:
211        move = restored.feed(token)
212        if move is not None:
213            recovered.append(move)
214    if recovered != moves or restored.finish().fen() != board.fen():
215        raise RuntimeError('round-trip mismatch')
216    return headers, tokens, ' '.join(m.uci() for m in moves), board.fen()
217
218
219def filter_reason(headers, config):
220    if config.get('completed') and headers['Result'] == '*':
221        return 'unfinished'
222    if config.get('exclude_bots') and any(v.upper() == 'BOT' for k, v in headers.items() if k in ('WhiteTitle', 'BlackTitle')):
223        return 'bot'
224    if config.get('min_rating') is not None:
225        for key in ('WhiteElo', 'BlackElo'):
226            if not re.fullmatch(r'[0-9]{1,10}', headers.get(key, '')):
227                return 'missing_or_nonnumeric_rating'
228            if int(headers[key]) < config['min_rating']:
229                return 'rating'
230    if config.get('min_time') is not None:
231        match = re.fullmatch(r'([0-9]{1,10})\+([0-9]{1,10})', headers.get('TimeControl', ''))
232        if not match:
233            return 'missing_or_unsupported_time_control'
234        if int(match[1]) + 40 * int(match[2]) < config['min_time']:
235            return 'time_control'
236    return None
237