Validate game headers (fixes niklasf#258)

niklasf · niklasf · commit c19acb9768b0 · 2018-01-11T15:23:48.000+01:00
diff --git a/chess/pgn.py b/chess/pgn.py
@@ -80,6 +80,8 @@
 
 TAG_REGEX = re.compile(r"^\[([A-Za-z0-9_]+)\s+\"(.*)\"\]\s*$")
 
+TAG_NAME_REGEX = re.compile(r"^[A-Za-z9-9_]+\Z")
+
 MOVETEXT_REGEX = re.compile(r"""
     (
         [NBKRQ]?[a-h]?[1-8]?[\-x]?[a-h][1-8](?:=?[nbrqkNBRQK])?
@@ -99,6 +101,9 @@
     """, re.DOTALL | re.VERBOSE)
 
 
+TAG_ROSTER = ["Event", "Site", "Date", "Round", "White", "Black", "Result"]
+
+
 class GameNode(object):
 
     def __init__(self):
@@ -389,18 +394,9 @@ class Game(GameNode):
     :class:`~chess.pgn.GameNode`.
     """
 
-    def __init__(self):
+    def __init__(self, headers=None):
         super(Game, self).__init__()
-
-        self.headers = collections.OrderedDict()
-        self.headers["Event"] = "?"
-        self.headers["Site"] = "?"
-        self.headers["Date"] = "????.??.??"
-        self.headers["Round"] = "?"
-        self.headers["White"] = "?"
-        self.headers["Black"] = "?"
-        self.headers["Result"] = "*"
-
+        self.headers = Headers(headers)
         self.errors = []
 
     def board(self, _cache=False):
@@ -504,11 +500,70 @@ def from_board(cls, board):
     @classmethod
     def without_tag_roster(cls):
         """Creates an empty game without the default 7 tag roster."""
-        game = cls.__new__(cls)
-        super(Game, game).__init__()
-        game.headers = collections.OrderedDict()
-        game.errors = []
-        return game
+        return cls(headers=Headers({}))
+
+
+class Headers(collections.MutableMapping):
+    def __init__(self, data=None, **kwargs):
+        self._tag_roster = {}
+        self._others = {}
+
+        if data is None:
+            data = {
+                "Event": "?",
+                "Site": "?",
+                "Date": "????.??.??",
+                "Round": "?",
+                "White": "?",
+                "Black": "?",
+                "Result": "*"
+            }
+
+        self.update(data, **kwargs)
+
+    def __setitem__(self, key, value):
+        if key in TAG_ROSTER:
+            self._tag_roster[key] = value
+        elif not TAG_NAME_REGEX.match(key):
+            raise ValueError("non-alphanumeric pgn header tag: {0}".format(repr(key)))
+        elif "\n" in value or "\r" in value:
+            raise ValueError("line break in pgn header {0}: {1}".format(key, repr(value)))
+        else:
+            self._others[key] = value
+
+    def __getitem__(self, key):
+        if key in TAG_ROSTER:
+            return self._tag_roster[key]
+        else:
+            return self._others[key]
+
+    def __delitem__(self, key):
+        if key in TAG_ROSTER:
+            del self._tag_roster[key]
+        else:
+            del self._others[key]
+
+    def __iter__(self):
+        for key in TAG_ROSTER:
+            if key in self._tag_roster:
+                yield key
+
+        for key in sorted(self._others):
+            yield key
+
+    def __len__(self):
+        return len(self._tag_roster) + len(self._others)
+
+    def copy(self):
+        return type(self)(self)
+
+    def __copy__(self):
+        return self.copy()
+
+    def __repr__(self):
+        return "{0}({1})".format(
+            type(self).__name__,
+            ", ".join("{0}={1}".format(key, repr(value)) for key, value in self.items()))
 
 
 class BaseVisitor(object):
@@ -1003,7 +1058,7 @@ def scan_headers(handle):
     Scan a PGN file opened in text mode for game offsets and headers.
 
     Yields a tuple for each game. The first element is the offset and the
-    second element is an ordered dictionary of game headers.
+    second element is a mapping of game headers.
 
     Since actually parsing many games from a big file is relatively expensive,
     this is a better way to look only for specific games and then seek and
@@ -1057,7 +1112,7 @@ def scan_headers(handle):
             tag_match = TAG_REGEX.match(line)
             if tag_match:
                 if game_pos is None:
-                    game_headers = Game().headers
+                    game_headers = Headers()
                     game_pos = last_pos
 
                 game_headers[tag_match.group(1)] = tag_match.group(2)
diff --git a/docs/pgn.rst b/docs/pgn.rst
@@ -52,8 +52,7 @@ holds general information, such as game headers.
 
     .. py:attribute:: headers
 
-        A :class:`collections.OrderedDict` of game headers. By default, the
-        following 7 headers are provided:
+        A mapping of headers. By default, the following 7 headers are provided:
 
         >>> import chess.pgn
         >>>
diff --git a/test.py b/test.py
@@ -2071,8 +2071,8 @@ def test_black_to_move(self):
             [White "?"]
             [Black "?"]
             [Result "*"]
-            [SetUp "1"]
             [FEN "8/8/4k3/8/4P3/4K3/8/8 b - - 0 17"]
+            [SetUp "1"]
 
             17... Kd6 18. Kd4 Ke6 *""")
 
